Fix OOM in test_large_block_sizes (#113153)
This test is causing flakyness on CPU, see #113134
Pull Request resolved: https://github.com/pytorch/pytorch/pull/113153
Approved by: https://github.com/lezcano
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
index 1e9ce1c..c6944ab 100644
--- a/test/inductor/test_torchinductor.py
+++ b/test/inductor/test_torchinductor.py
@@ -7481,9 +7481,11 @@
Inductor will try triton configs like x = 64 and y = 1024 which will
result in out of shared memory if dtype is fp32.
- Currnelty inductor will skip such bad configs and pick the best one
+ Currently inductor will skip such bad configs and pick the best one
from the remaining configs.
"""
+ if not _has_sufficient_memory(self.device, 3 * 2**24 * 65 * 4):
+ raise unittest.SkipTest("insufficient memory")
@torch.compile
def fn(x, y):
@@ -7491,11 +7493,8 @@
# Use shape (2**24, 65) rather than (2**24, 128) potentially avoid OOM in
# CI while still keep the same up-rounded size-hints.
- try:
- a = torch.randn(2**24, 65, device=self.device)
- b = torch.randn(65, 2**24, device=self.device)
- except RuntimeError:
- return # skip testing if OOM
+ a = torch.randn(2**24, 65, device=self.device)
+ b = torch.randn(65, 2**24, device=self.device)
fn(a, b)
