diff --git a/test/test_cuda.py b/test/test_cuda.py
index 35e834d3b62e..bde798913cf4 100644
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -4305,7 +4305,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):