update

Lightning-Universe · Apr 19, 2021 · 86eaaaf · 86eaaaf
1 parent a06829a
commit 86eaaaf
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/data/test_data_pipeline.py b/tests/data/test_data_pipeline.py
@@ -598,7 +598,7 @@ def val_per_batch_transform_on_device(self, batch: Any) -> Any:
             batch = batch[0]
         assert torch.equal(batch["a"], tensor([0, 1]))
         assert torch.equal(batch["b"], tensor([1, 2]))
-        return False
+        return [False]
 
     @staticmethod
     def fn_test_load_data() -> List[torch.Tensor]:
@@ -649,6 +649,8 @@ def training_step(self, batch, batch_idx):
         assert batch is None
 
     def validation_step(self, batch, batch_idx):
+        if isinstance(batch, list):
+            batch = batch[0]
         assert batch is False
 
     def test_step(self, batch, batch_idx):