Fixes bugs in stitcher.py that show up in distributed setting

neuro-galaxy · Nov 11, 2024 · bf14232 · bf14232
1 parent 1bc9c73
commit bf14232
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -22,6 +22,7 @@
         "wandb~=0.15",
         "torchtyping~=0.1",
         "pydantic~=2.0",
+        "tabulate",
     ],
     extras_require={
         "dev": [

diff --git a/torch_brain/utils/stitcher.py b/torch_brain/utils/stitcher.py
@@ -189,16 +189,21 @@ def on_validation_epoch_end(self, trainer, pl_module, prefix="val"):
             for task_name in self.metrics[recording_id].keys():
                 for metric_name in self.metrics[recording_id][task_name].keys():
                     metrics[f"{recording_id}/{task_name}/{metric_name}/{prefix}"] = (
-                        self.metrics[recording_id][task_name][metric_name].compute()
+                        self.metrics[recording_id][task_name][metric_name]
+                        .to(pl_module.device)
+                        .compute()
                     )
                     self.metrics[recording_id][task_name][metric_name].reset()
+                    self.metrics[recording_id][task_name][metric_name].to("cpu")
 
         # log the metrics
         self.log_dict(metrics)
         logging.info(f"Logged {len(metrics)} {prefix} metrics.")
 
         # compute the average metric
-        metrics[f"average_{prefix}_metric"] = np.array(list(metrics.values())).mean()
+        metrics[f"average_{prefix}_metric"] = torch.tensor(
+            list(metrics.values())
+        ).mean()
 
         metrics_data = []
         for metric_name, metric_value in metrics.items():