diff --git a/test/test_cuda.py b/test/test_cuda.py
index c707433ebe6e4..b6b1cd07cb22d 100644
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -4475,7 +4475,7 @@ def alloc(n):
             # expandable_segment blocks can be in the free list when this is called.
             alloc(80)
         finally:
-            orig = torch.cuda.get_per_process_memory_fraction(0)
+            torch.cuda.memory.set_per_process_memory_fraction(orig)
 
     def test_allocator_settings(self):
         def power2_div(size, div_factor):