From 3610b76389ab478a54db87196f1c9b2e1d9c4c5f Mon Sep 17 00:00:00 2001 From: benjijamorris <54606172+benjijamorris@users.noreply.github.com> Date: Wed, 14 Feb 2024 16:15:09 -0800 Subject: [PATCH] add filename to batch for testing (#333) Co-authored-by: Benjamin Morris --- cyto_dl/models/im2im/multi_task.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cyto_dl/models/im2im/multi_task.py b/cyto_dl/models/im2im/multi_task.py index 6fc08e3fe..ada16b853 100644 --- 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):