fix test_set_per_process_memory_fraction failure caused by knock-on effect of test_out_of_memory_retry

mcarilli · mcarilli · commit c85cb9c779f4 · 2022-02-09T18:15:09.000-07:00
diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp
@@ -91,8 +91,6 @@ namespace Native {
  *                  notifyCaptureDestroy.
  */
 
-namespace {
-
 constexpr size_t kMinBlockSize =
     512; // all sizes are rounded to at least 512 bytes
 constexpr size_t kSmallSize = 1048576; // largest "small" allocation is 1 MiB
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -378,6 +378,7 @@ def test_out_of_memory(self):
         self.assertTrue((tensor == 1).all())
 
     def test_out_of_memory_retry(self):
+        torch.cuda.empty_cache()
         total_memory = torch.cuda.get_device_properties(0).total_memory
         oom_regex = "would exceed allowed memory" if TEST_CUDAMALLOCASYNC else \
                     "Tried to allocate"
@@ -387,6 +388,10 @@ def test_out_of_memory_retry(self):
             b = torch.empty(size, dtype=torch.int8, device='cuda')
         del a
         b = torch.empty(size, dtype=torch.int8, device='cuda')
+        del b
+        # We used a lot of memory here, clean up so we don't affect other tests too much
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
 
     def test_set_per_process_memory_fraction(self):
         # test invalid fraction value.
