test_complex inherits from common_utils.TestCase; closes #34648 (#34697)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/34697
Differential Revision: D20596860
Pulled By: anjali411
fbshipit-source-id: 18599fce5bd3513be17ecf83ba9fb0d64d971fc4
diff --git a/test/test_complex.py b/test/test_complex.py
index 9d6fbaf..3912fe7 100644
--- a/test/test_complex.py
+++ b/test/test_complex.py
@@ -1,13 +1,13 @@
import torch
-import unittest
+from torch.testing._internal.common_utils import TestCase, run_tests
devices = (torch.device('cpu'), torch.device('cuda:0'))
-class TestComplexTensor(unittest.TestCase):
+class TestComplexTensor(TestCase):
def test_to_list_with_complex_64(self):
# test that the complex float tensor has expected values and
# there's no garbage value in the resultant list
self.assertEqual(torch.zeros((2, 2), dtype=torch.complex64).tolist(), [[0j, 0j], [0j, 0j]])
if __name__ == '__main__':
- unittest.main()
+ run_tests()
