Split CUDA and CPU tests of torchscript_consistency

pytorch · May 15, 2020 · cb5eacd · cb5eacd
1 parent 2c28b74
commit cb5eacd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/test/torchscript_consistency_cpu_test.py b/test/torchscript_consistency_cpu_test.py
@@ -0,0 +1,5 @@
+from common_utils import define_test_suites
+from torchscript_consistency_impl import Functional, Transforms
+
+
+define_test_suites(globals(), [Functional, Transforms], devices=['cpu'])
diff --git a/test/torchscript_consistency_cuda_test.py b/test/torchscript_consistency_cuda_test.py
@@ -0,0 +1,5 @@
+from common_utils import define_test_suites
+from torchscript_consistency_impl import Functional, Transforms
+
+
+define_test_suites(globals(), [Functional, Transforms], devices=['cuda'])
diff --git a/test/test_torchscript_consistency.py → test/torchscript_consistency_impl.py b/test/test_torchscript_consistency.py → test/torchscript_consistency_impl.py
@@ -608,6 +608,3 @@ def test_Vad(self):
         filepath = common_utils.get_asset_path("vad-go-mono-32000.wav")
         waveform, sample_rate = torchaudio.load(filepath)
         self._assert_consistency(T.Vad(sample_rate=sample_rate), waveform)
-
-
-common_utils.define_test_suites(globals(), [Functional, Transforms])