Skip to content

Commit cf9cad3

Browse files
AnthonyBarbierpytorchmergebot
authored andcommitted
Add __main__ guards to tests (pytorch#154716)
This PR is part of a series attempting to re-submit pytorch#134592 as smaller PRs. Add missing `if __name__ == "__main__":` guards to some tests. Pull Request resolved: pytorch#154716 Approved by: https://github.com/Skylion007
1 parent ca0c298 commit cf9cad3

9 files changed

+80
-17
lines changed

test/inductor/test_aot_inductor_utils.py

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from torch._inductor import config
1717
from torch._inductor.test_case import TestCase
1818
from torch.testing import FileCheck
19-
from torch.testing._internal.common_utils import IS_FBCODE
19+
from torch.testing._internal.common_utils import IS_FBCODE, run_tests
2020
from torch.testing._internal.inductor_utils import clone_preserve_strides_offset
2121
from torch.utils import _pytree as pytree
2222

@@ -205,11 +205,14 @@ def check_model(
205205
atol=None,
206206
rtol=None,
207207
):
208-
with torch.no_grad(), config.patch(
209-
{
210-
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
211-
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
212-
}
208+
with (
209+
torch.no_grad(),
210+
config.patch(
211+
{
212+
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
213+
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
214+
}
215+
),
213216
):
214217
torch.manual_seed(0)
215218
if not isinstance(model, types.FunctionType):
@@ -248,11 +251,14 @@ def check_model_with_multiple_inputs(
248251
options=None,
249252
dynamic_shapes=None,
250253
):
251-
with torch.no_grad(), config.patch(
252-
{
253-
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
254-
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
255-
}
254+
with (
255+
torch.no_grad(),
256+
config.patch(
257+
{
258+
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
259+
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
260+
}
261+
),
256262
):
257263
torch.manual_seed(0)
258264
model = model.to(self.device)
@@ -275,11 +281,14 @@ def code_check_count(
275281
target_str: str,
276282
target_count: int,
277283
):
278-
with torch.no_grad(), config.patch(
279-
{
280-
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
281-
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
282-
}
284+
with (
285+
torch.no_grad(),
286+
config.patch(
287+
{
288+
"aot_inductor.allow_stack_allocation": self.allow_stack_allocation,
289+
"aot_inductor.use_minimal_arrayref_interface": self.use_minimal_arrayref_interface,
290+
}
291+
),
283292
):
284293
package_path = torch._export.aot_compile(model, example_inputs)
285294

@@ -290,3 +299,7 @@ def code_check_count(
290299
target_count,
291300
exactly=True,
292301
).run(src_code)
302+
303+
304+
if __name__ == "__main__":
305+
run_tests()

test/lazy/test_bindings.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
11
# Owner(s): ["oncall: jit"]
22

33
import torch._lazy.metrics
4+
from torch.testing._internal.common_utils import run_tests
45

56

67
def test_metrics():
78
names = torch._lazy.metrics.counter_names()
89
assert len(names) == 0, f"Expected no counter names, but got {names}"
10+
11+
12+
if __name__ == "__main__":
13+
run_tests()

test/lazy/test_extract_compiled_graph.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,3 +206,10 @@ class OptimizeTest(unittest.TestCase):
206206
test_return_multi = maketest(ModuleReturnMulti)
207207
test_return_dup_tensor = maketest(ModuleReturnDupTensor)
208208
test_inplace_update = maketest(ModuleInplaceUpdate)
209+
210+
211+
if __name__ == "__main__":
212+
raise RuntimeError(
213+
"This test is not currently used and should be "
214+
"enabled in discover_tests.py if required."
215+
)

test/lazy/test_meta_kernel.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,10 @@ def test_addmm(self):
3737
def test_add_invalid_device(self):
3838
with self.assertRaisesRegex(RuntimeError, ".*not a lazy tensor.*"):
3939
_ = torch.tensor([1], device="cpu") + torch.tensor([1], device="lazy")
40+
41+
42+
if __name__ == "__main__":
43+
raise RuntimeError(
44+
"This test is not currently used and should be "
45+
"enabled in discover_tests.py if required."
46+
)

test/onnx/test_models_quantized_onnxruntime.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import torch
1212
from torch import nn
13+
from torch.testing._internal import common_utils
1314

1415

1516
def _get_test_image_tensor():
@@ -95,3 +96,7 @@ def test_resnext101_32x8d(self):
9596
pretrained=True, quantize=True
9697
)
9798
self.run_test(model, _get_test_image_tensor())
99+
100+
101+
if __name__ == "__main__":
102+
common_utils.run_tests()

test/onnx/test_onnxscript_no_runtime.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,3 +160,10 @@ def custom_selu(g, X):
160160
)
161161
loop_selu_proto = onnx.load(io.BytesIO(saved_model.getvalue()))
162162
self.assertEqual(len(loop_selu_proto.functions), 1)
163+
164+
165+
if __name__ == "__main__":
166+
raise RuntimeError(
167+
"This test is not currently used and should be "
168+
"enabled in discover_tests.py if required."
169+
)

test/test_bundled_images.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,10 @@ def forward(self, arg):
9292
im2_tensor = torch.ops.fb.image_decode_to_NCHW(byte_tensor, weight, bias)
9393
self.assertEqual(raw_data.shape, im2_tensor.shape)
9494
self.assertEqual(raw_data, im2_tensor, atol=0.1, rtol=1e-01)
95+
96+
97+
if __name__ == "__main__":
98+
raise RuntimeError(
99+
"This test is not currently used and should be "
100+
"enabled in discover_tests.py if required."
101+
)

test/test_hub.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,12 @@
88

99
import torch
1010
import torch.hub as hub
11-
from torch.testing._internal.common_utils import IS_SANDCASTLE, retry, TestCase
11+
from torch.testing._internal.common_utils import (
12+
IS_SANDCASTLE,
13+
retry,
14+
run_tests,
15+
TestCase,
16+
)
1217

1318

1419
def sum_of_state_dict(state_dict):
@@ -307,3 +312,7 @@ def test_trust_repo_legacy(self):
307312
torch.hub.load("ailzhang/torchhub_example", "mnist_zip_1_6", trust_repo="check")
308313

309314
self._assert_trusted_list_is_empty()
315+
316+
317+
if __name__ == "__main__":
318+
run_tests()

tools/testing/discover_tests.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,10 @@ def skip_test_p(name: str) -> bool:
104104
"distributed/test_c10d_spawn",
105105
"distributions/test_transforms",
106106
"distributions/test_utils",
107+
"lazy/test_meta_kernel",
108+
"lazy/test_extract_compiled_graph",
107109
"test/inductor/test_aot_inductor_utils",
110+
"onnx/test_onnxscript_no_runtime",
108111
"onnx/test_pytorch_onnx_onnxruntime_cuda",
109112
"onnx/test_models",
110113
# These are not C++ tests

0 commit comments

Comments
 (0)