chore: Fix dynamo tests

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/dynamo/test/trt_lower/test_fx2trt_lower.py

@@ -20,7 +20,7 @@ def forward(self, x):
 
         mod = _Mod()
         mod_traced = fx.symbolic_trace(mod)
-        input = [torch.rand(4)]
+        input = [torch.rand(4).cuda()]
         lower = Lowerer.create(LowerSetting())
         lower(mod_traced, input)
 
@@ -39,7 +39,7 @@ def forward(self, x):
                 return self.bn(x)
 
         module = TestModule()
-        inputs = [torch.randn(1, 3, 224, 224)]
+        inputs = [torch.randn(1, 3, 224, 224).cuda()]
         lower = Lowerer.create(LowerSetting(ast_rewriter_allow_list={MyBatchNorm}))
         lower(module, inputs)
 
@@ -53,7 +53,7 @@ def forward(self, x):
                 return (torch.sqrt(x), self.a)
 
         lower = Lowerer.create(LowerSetting())
-        lower(TestModule(), [torch.randn([2, 2])])
+        lower(TestModule(), [torch.randn([2, 2]).cuda()])
 
     def test_replace_mutable_op(self):
         class TestModule(torch.nn.Module):
@@ -65,7 +65,7 @@ def forward(self, x, y):
 
         lower = Lowerer.create(LowerSetting())
         mod_traced = fx.symbolic_trace(TestModule())
-        lower(mod_traced, [torch.randn(3, 4), torch.randn(3, 4)])
+        lower(mod_traced, [torch.randn(3, 4).cuda(), torch.randn(3, 4).cuda()])
 
     def test_replace_mutable_op_dont_apply(self):
         class TestModule(torch.nn.Module):

py/torch_tensorrt/dynamo/test/trt_lower/test_observer_gpu.py

@@ -19,8 +19,8 @@ def test_observe_lowerer(self):
         import torch
         import torch.nn as nn
 
-        import torch_tensorrt.fx.lower as lower
-        from torch_tensorrt.fx.lower_setting import LowerSetting
+        import torch_tensorrt.dynamo.lower as lower
+        from torch_tensorrt.dynamo.lower_setting import LowerSetting
 
         class Model(nn.Module):
             def forward(self, x, y):

py/torch_tensorrt/dynamo/test/trt_lower/trt_splitter_test.py

@@ -358,12 +358,12 @@ def test_splitter(splitter):
 
         test_splitter(splitter)
 
-    def test_min_block_size(self):
+    def test_min_acc_module_size(self):
         """
            sin     relu     cos     sigmoid     tanh
         a ====> b =====> c ====> d ========> e =====> f
 
-        We set sin, cos and tanh as acc node but also set min_block_size to 2
+        We set sin, cos and tanh as acc node but also set min_acc_module_size to 2
         and expect the whole module stay on CPU.
         """
 
@@ -386,9 +386,9 @@ class CustomOpSupport(op_support.OperatorSupport):
                 "acc_ops.tanh": None,
             }
 
-        # Create splitter setting and set min_block_size to 2
+        # Create splitter setting and set min_acc_module_size to 2
         settings = splitter_base._SplitterSettingBase()
-        settings.min_block_size = 2
+        settings.min_acc_module_size = 2
         splitter = TRTSplitter(
             mod,
             (torch.randn(2, 3),),
@@ -815,7 +815,7 @@ def test_split_non_tensor_edges_2(self):
         # Making 'a', 'b1', 'b2', 'd' and 'e' run on ACC with limit on ACC
         # subgraph size
         settings = splitter_base._SplitterSettingBase()
-        settings.min_block_size = 2
+        settings.min_acc_module_size = 2
         splitter = TRTSplitter(
             module_nn,
             (test_data,),
@@ -912,7 +912,7 @@ def test_split_non_tensor_edges_4(self):
         # Making 'a', 'c', 'd' and 'e' run on ACC with limit on ACC
         # subgraph size
         settings = splitter_base._SplitterSettingBase()
-        settings.min_block_size = 2
+        settings.min_acc_module_size = 2
         splitter = TRTSplitter(
             module_nn,
             (test_data,),
@@ -1072,7 +1072,7 @@ def test_start_with_acc_module_(self):
            sin     relu     cos     sigmoid     tanh
         a ====> b =====> c ====> d ========> e =====> f
 
-        We set sin, relu and cos as acc node but also set min_block_size to 2
+        We set sin, relu and cos as acc node but also set min_acc_module_size to 2
         and expect the whole module stay on CPU.
         """
 
@@ -1095,9 +1095,9 @@ class CustomOpSupport(op_support.OperatorSupport):
                 "acc_ops.relu": None,
             }
 
-        # Create splitter setting and set min_block_size to 2
+        # Create splitter setting and set min_acc_module_size to 2
         settings = splitter_base._SplitterSettingBase()
-        settings.min_block_size = 2
+        settings.min_acc_module_size = 2
         splitter = TRTSplitter(
             mod,
             (torch.randn(2, 3),),