From cf8f1ffed306af7f99eeee5ef9452bada69827ab Mon Sep 17 00:00:00 2001 From: Andrey Bokovoy Date: Thu, 9 Apr 2026 11:46:30 +0000 Subject: [PATCH] Get actual free GPU memory in test_cache_int32_overflow --- fbgemm_gpu/test/tbe/cache/cache_overflow_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fbgemm_gpu/test/tbe/cache/cache_overflow_test.py b/fbgemm_gpu/test/tbe/cache/cache_overflow_test.py index b4c4e5b764..8da8b199f1 100644 --- a/fbgemm_gpu/test/tbe/cache/cache_overflow_test.py +++ b/fbgemm_gpu/test/tbe/cache/cache_overflow_test.py @@ -39,8 +39,8 @@ def test_cache_int32_overflow(self, stochastic_rounding: bool) -> None: cache_sets = 10**6 current_device = torch.device(torch.cuda.current_device()) - total_memory = torch.cuda.get_device_properties(current_device).total_memory - free_memory = total_memory - torch.cuda.memory_reserved(current_device) + free_memory, _ = torch.cuda.mem_get_info(current_device) + free_memory = int(free_memory * 0.8) # Weight and cache precisions are fixed to FP16 element_size = 2