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):