check which assert is failing

Lightning-AI · Jun 17, 2021 · b7483b4 · b7483b4
1 parent cab35a6
commit b7483b4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/trainer/loops/test_evaluation_loop.py b/tests/trainer/loops/test_evaluation_loop.py
@@ -95,14 +95,16 @@ def training_step(self, batch, batch_idx):
             # there is a batch and the boring model, but not two batches on gpu, assume 32 bit = 4 bytes
             lower = 101 * self.num_params * 4
             upper = 201 * self.num_params * 4
-            assert lower < torch.cuda.memory_allocated(0) < upper
+            assert lower < torch.cuda.memory_allocated(self.device)
+            assert torch.cuda.memory_allocated(self.device) < upper
             return super().training_step(batch, batch_idx)
 
         def validation_step(self, batch, batch_idx):
             # there is a batch and the boring model, but not two batches on gpu, assume 32 bit = 4 bytes
             lower = 101 * self.num_params * 4
             upper = 201 * self.num_params * 4
-            assert lower < torch.cuda.memory_allocated(0) < upper
+            assert lower < torch.cuda.memory_allocated(self.device)
+            assert torch.cuda.memory_allocated(self.device) < upper
             return super().validation_step(batch, batch_idx)
 
     torch.cuda.empty_cache()