[inductor] Fix flaky tests in test_benchmark_fusion.py (#131733)

Summary: Same as [#131699](https://github.com/pytorch/pytorch/pull/131699), but in `test_benchmark_fusion.py`.

Test Plan: Tested internally.

Differential Revision: D60211793

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131733
Approved by: https://github.com/oulgen
This commit is contained in:
Adnan Akhundov
2024-07-25 05:39:14 +00:00
committed by PyTorch MergeBot
parent 42a4df9447
commit ee996cd63c

View File

@ -5,7 +5,7 @@ import sys
import torch
from torch._inductor.test_case import TestCase as InductorTestCase
from torch._inductor.utils import fresh_inductor_cache, run_and_get_code
from torch._inductor.utils import fresh_inductor_cache, is_big_gpu, run_and_get_code
from torch.testing import FileCheck
from torch.testing._internal.common_utils import (
IS_CI,
@ -211,6 +211,11 @@ if HAS_CUDA and not TEST_WITH_ASAN:
cls._stack.close()
super().tearDownClass()
def setUp(self):
super().setUp()
if not is_big_gpu(0):
return self.skipTest("Need a big GPU to run max_autotune=True")
def _equivalent_output_code_impl(self, size, first_dim=None, activation=True):
def foo(m, inp):
a = m(inp)