Add comments to the CPU memory test

BlackHC · Apr 17, 2021 · 10cfe70 · 10cfe70
1 parent 84a848e
commit 10cfe70
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_cpu_mem_limit.py b/tests/test_cpu_mem_limit.py
@@ -6,10 +6,12 @@
 
 @pytest.mark.forked
 def test_cpu_mem_limit():
+    # 128 MB (128/4M float32)
     tensor = torch.empty((128, 1024, 1024 // 4), dtype=torch.float32)
     tensor.resize_(1)
 
     cpu_memory.set_cpu_memory_limit(0.25)
 
+    # 512 MB (128/4M float32)
     with pytest.raises(RuntimeError):
         torch.empty((512, 1024, 1024 // 4), dtype=torch.float32)