Added test case

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

tests/py/multi_gpu_test_case.py

@@ -0,0 +1,21 @@
+import unittest
+import trtorch
+import torch
+import torchvision.models as models
+
+
+class MultiGpuTestCase(unittest.TestCase):
+
+    def __init__(self, methodName='runTest', model=None):
+        super(MultiGpuTestCase, self).__init__(methodName)
+        self.model = model
+        self.model.eval().to("cuda")
+
+    @staticmethod
+    def parametrize(testcase_class, model=None):
+        testloader = unittest.TestLoader()
+        testnames = testloader.getTestCaseNames(testcase_class)
+        suite = unittest.TestSuite()
+        for name in testnames:
+            suite.addTest(testcase_class(name, model=model))
+        return suite

tests/py/test_api_multi_gpu.py

@@ -3,17 +3,12 @@
 import torch
 import torchvision.models as models
 
-from model_test_case import ModelTestCase
+from multi_gpu_test_case import MultiGpuTestCase
 
+gpu_id = 1
 class TestCompile(MultiGpuTestCase):
 
     def setUp(self):
-        if not torch.cuda.device_count() > 1:
-            raise ValueError("This test case is applicable for multi-gpu configurations only")
-        
-        self.gpu_id = 1
-        # Setting it up here so that all CUDA allocations are done on correct device
-        trtorch.set_device(self.gpu_id)
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
         self.traced_model = torch.jit.trace(self.model, [self.input])
         self.scripted_model = torch.jit.script(self.model)
@@ -23,7 +18,7 @@ def test_compile_traced(self):
             "input_shapes": [self.input.shape],
             "device": {
                 "device_type": trtorch.DeviceType.GPU,
-                "gpu_id": self.gpu_id,
+                "gpu_id": gpu_id,
                 "dla_core": 0,
                 "allow_gpu_fallback": False,
                 "disable_tf32": False
@@ -39,7 +34,7 @@ def test_compile_script(self):
             "input_shapes": [self.input.shape],
             "device": {
                 "device_type": trtorch.DeviceType.GPU,
-                "gpu_id": self.gpu_id,
+                "gpu_id": gpu_id,
                 "dla_core": 0,
                 "allow_gpu_fallback": False,
                 "disable_tf32": False
@@ -58,6 +53,11 @@ def test_suite():
 
     return suite
 
+if not torch.cuda.device_count() > 1:
+    raise ValueError("This test case is applicable for multi-gpu configurations only")
+        
+# Setting it up here so that all CUDA allocations are done on correct device
+trtorch.set_device(gpu_id)
 suite = test_suite()
 
 runner = unittest.TextTestRunner()