Add missing __main__ in two unittests (pytorch#97302)

Pull Request resolved: pytorch#97302 Approved by: https://github.com/zou3519
walterddr · Mar 22, 2023 · 726fc36 · 726fc36
1 parent 28929b1
commit 726fc36
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/test/test_comparison_utils.py b/test/test_comparison_utils.py
@@ -2,7 +2,7 @@
 # Owner(s): ["module: internals"]
 
 import torch
-from torch.testing._internal.common_utils import TestCase
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 class TestComparisonUtils(TestCase):
     def test_all_equal_no_assert(self):
@@ -30,3 +30,7 @@ def test_assert_sizes(self):
 
         with self.assertRaises(RuntimeError):
             torch._assert_tensor_metadata(t, [3], [1], torch.float)
+
+
+if __name__ == '__main__':
+    run_tests()
diff --git a/test/test_pruning_op.py b/test/test_pruning_op.py
@@ -4,7 +4,7 @@
 from hypothesis import given
 import numpy as np
 import torch
-from torch.testing._internal.common_utils import TestCase
+from torch.testing._internal.common_utils import TestCase, run_tests
 import torch.testing._internal.hypothesis_utils as hu
 hu.assert_deadline_disabled()
 
@@ -76,3 +76,7 @@ def test_rowwise_prune_op_32bit_indices(self, embedding_rows, embedding_dims, we
     )
     def test_rowwise_prune_op_64bit_indices(self, embedding_rows, embedding_dims, weights_dtype):
         self._test_rowwise_prune_op(embedding_rows, embedding_dims, torch.int64, weights_dtype)
+
+
+if __name__ == '__main__':
+    run_tests()