pig

Lightning-AI · Jun 17, 2021 · 33d89e0 · 33d89e0
1 parent b7483b4
commit 33d89e0
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/tests/trainer/loops/test_evaluation_loop.py b/tests/trainer/loops/test_evaluation_loop.py
@@ -77,6 +77,9 @@ def on_validation_end(self):
 def test_memory_consumption_validation(tmpdir):
     """Test that the training batch is no longer in GPU memory when running validation"""
 
+    memory_start = torch.cuda.memory_allocated(0)
+    print("horse memory: ", memory_start)
+
     class BoringLargeBatchModel(BoringModel):
 
         @property
@@ -93,18 +96,19 @@ def val_dataloader(self):
 
         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
+            print("pig memory: ", memory_start)
             lower = 101 * self.num_params * 4
             upper = 201 * self.num_params * 4
-            assert lower < torch.cuda.memory_allocated(self.device)
-            assert torch.cuda.memory_allocated(self.device) < upper
+            assert lower < torch.cuda.memory_allocated(self.device) - memory_start
+            assert torch.cuda.memory_allocated(self.device) - memory_start < 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(self.device)
-            assert torch.cuda.memory_allocated(self.device) < upper
+            assert lower < torch.cuda.memory_allocated(self.device) - memory_start
+            assert torch.cuda.memory_allocated(self.device) - memory_start < upper
             return super().validation_step(batch, batch_idx)
 
     torch.cuda.empty_cache()