add filename to batch for testing (#333)

Co-authored-by: Benjamin Morris <[email protected]>
AllenCellModeling · Feb 15, 2024 · 3610b76 · 3610b76
1 parent 87c7617
commit 3610b76
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/cyto_dl/models/im2im/multi_task.py b/cyto_dl/models/im2im/multi_task.py
@@ -167,6 +167,7 @@ def _get_run_heads(self, batch, stage):
         return run_heads
 
     def model_step(self, stage, batch, batch_idx):
+        batch["filenames"] = batch[self.hparams.x_key].meta.get("filename_or_obj", batch_idx)
         # convert monai metatensors to tensors
         for k, v in batch.items():
             if isinstance(v, MetaTensor):