diff --git a/test/test_cuda.py b/test/test_cuda.py index b33f21e6dfce2..b00fb9934fed9 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -5045,12 +5045,16 @@ def test_throw_on_cudamalloc_oom(self): # preemptively rejected with OutOfMemoryError. # Both settings must go through _accelerator_setAllocatorSettings so # they are read from CUDAAllocatorConfig. + fraction = 0.005 torch._C._accelerator_setAllocatorSettings( - "throw_on_cudamalloc_oom:True,per_process_memory_fraction:0.01" + f"throw_on_cudamalloc_oom:True,per_process_memory_fraction:{fraction}" ) + total_mem = torch.cuda.get_device_properties(0).total_memory + # Allocate the allowed threshold + 1 MiB to guarantee rejection + alloc_bytes = int(total_mem * fraction) + 1024 * 1024 with self.assertRaises(torch.cuda.OutOfMemoryError): - torch.empty(1024 * 1024 * 1024, dtype=torch.int8, device="cuda") + torch.empty(alloc_bytes, dtype=torch.int8, device="cuda") # Check that rejection counter was incremented stats = torch.cuda.memory_stats() diff --git a/test/test_torch.py b/test/test_torch.py index 48a463a0d2959..e6991c0513fdd 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -828,6 +828,7 @@ def test_cpp_warnings_have_python_context(self, device): s = ".+Triggered internally at.+RangeFactories.+" # nvfuser deprecation warning filter warnings.filterwarnings("ignore", "torch::jit::fuser::cuda", UserWarning) + warnings.filterwarnings("ignore", ".*", DeprecationWarning) # ignore all deprecation warnings def cpp_warn_fn(): out = torch.empty((5,))