Set default value for dataloader_idx

ruflab · Nov 8, 2023 · 7a22542 · 7a22542
1 parent 1ef896b
commit 7a22542
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/shimmer/modules/global_workspace.py b/shimmer/modules/global_workspace.py
@@ -187,7 +187,7 @@ def generic_step(
         return losses["loss"]
 
     def validation_step(
-        self, data: Mapping[str, Any], _, dataloader_idx: int
+        self, data: Mapping[str, Any], _, dataloader_idx: int = 0
     ) -> torch.Tensor:
         batch = {frozenset(data.keys()): data}
         for domain in data.keys():
@@ -197,7 +197,7 @@ def validation_step(
         return self.generic_step(batch, mode="val/ood")
 
     def test_step(
-        self, data: Mapping[str, Any], _, dataloader_idx: int
+        self, data: Mapping[str, Any], _, dataloader_idx: int = 0
     ) -> torch.Tensor:
         batch = {frozenset(data.keys()): data}
         for domain in data.keys():