Patch ModelSummary with BatchTransforms (#307)

* adding the fix to run model summary

* remove redundant branch

* Revert "remove redundant branch"

This reverts commit 1a06cc1.

* remove fixme comment

---------

Co-authored-by: Ziwen Liu <[email protected]>

Author: edyoshikun

viscy/data/combined.py

@@ -231,15 +231,22 @@ def val_dataloader(self):
 
     def on_after_batch_transfer(self, batch, dataloader_idx: int):
         """Apply GPU transforms from constituent data modules to micro-batches."""
+        if not isinstance(batch, list):
+            return batch
+
         processed_micro_batches = []
         for micro_batch in batch:
-            dataset_idx = micro_batch.pop("_dataset_idx")
-            dm = self.data_modules[dataset_idx]
-            if hasattr(dm, "on_after_batch_transfer"):
-                processed_micro_batch = dm.on_after_batch_transfer(
-                    micro_batch, dataloader_idx
-                )
+            if isinstance(micro_batch, dict) and "_dataset_idx" in micro_batch:
+                dataset_idx = micro_batch.pop("_dataset_idx")
+                dm = self.data_modules[dataset_idx]
+                if hasattr(dm, "on_after_batch_transfer"):
+                    processed_micro_batch = dm.on_after_batch_transfer(
+                        micro_batch, dataloader_idx
+                    )
+                else:
+                    processed_micro_batch = micro_batch
             else:
+                # Handle case where micro_batch doesn't have _dataset_idx (e.g., from model summary)
                 processed_micro_batch = micro_batch
             processed_micro_batches.append(processed_micro_batch)
         combined_batch = {}