examples/research_projects/rag-end2end-retriever/finetune_rag.py/GenerativeQAModule/validation_epoch_end
class GenerativeQAModule:    def validation_epoch_end(self, outputs, prefix="val") -> Dict:
        self.step_count += 1
        losses = {k: torch.stack([x[k] for x in outputs]).mean() for k in self.loss_names}
        loss = losses["loss"]
        gen_metrics = {
            k: np.array([x[k] for x in outputs]).mean() for k in self.metric_names + ["gen_time", "gen_len"]
        }
        metrics_tensor: torch.FloatTensor = torch.tensor(gen_metrics[self.val_metric]).type_as(loss)
        gen_metrics.update({k: v.item() for k, v in losses.items()})

        # fix for https://github.com/PyTorchLightning/pytorch-lightning/issues/2424
        if dist.is_initialized():
            dist.all_reduce(metrics_tensor, op=dist.ReduceOp.SUM)
            metrics_tensor = metrics_tensor / dist.get_world_size()
            gen_metrics.update({self.val_metric: metrics_tensor.item()})

        losses.update(gen_metrics)
        metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
        metrics["step_count"] = self.step_count
        self.save_metrics(metrics, prefix)  # writes to self.metrics_save_path

        log_dict = {
            f"{prefix}_avg_em": metrics[f"{prefix}_avg_em"],
            "step_count": metrics["step_count"],
            f"{prefix}_avg_loss": metrics[f"{prefix}_avg_loss"],
            f"{prefix}_loss": loss,
            f"{prefix}_em": metrics_tensor,
        }
        self.log_dict(log_dict)