Fixed CI issues

Author: cehongwang

tests/py/dynamo/models/test_export_serde.py

@@ -52,6 +52,7 @@ def forward(self, x):
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     # Check Pyt and TRT exported program outputs
+    model.cuda()
     cos_sim = cosine_similarity(model(input), trt_module(input)[0])
     assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
@@ -106,6 +107,7 @@ def forward(self, x):
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     # Check Pyt and TRT exported program outputs
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
     for idx in range(len(outputs_pyt)):
@@ -162,8 +164,9 @@ def forward(self, x):
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
     torchtrt.save(trt_module, trt_ep_path)
 
-    deser_trt_module = torchtrt.load(trt_ep_path).module()
+    deser_trt_module = torchtrt.load(trt_ep_path).module().cuda()
     # Check Pyt and TRT exported program outputs
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
     for idx in range(len(outputs_pyt)):
@@ -173,7 +176,7 @@ def forward(self, x):
             msg=f"test_no_compile TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
         )
 
-    # # Check Pyt and deserialized TRT exported program outputs
+    # Check Pyt and deserialized TRT exported program outputs
     outputs_trt_deser = deser_trt_module(input)
     for idx in range(len(outputs_pyt)):
         cos_sim = cosine_similarity(outputs_pyt[idx], outputs_trt_deser[idx])
@@ -224,6 +227,7 @@ def forward(self, x):
     torchtrt.save(trt_module, trt_ep_path)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
     for idx in range(len(outputs_pyt)):
@@ -267,6 +271,7 @@ def test_resnet18(ir):
     torchtrt.save(trt_module, trt_ep_path)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
     cos_sim = cosine_similarity(outputs_pyt, outputs_trt[0])
@@ -312,6 +317,7 @@ def test_resnet18_dynamic(ir):
     torchtrt.save(trt_module, trt_ep_path)
     # TODO: Enable this serialization issues are fixed
     # deser_trt_module = torchtrt.load(trt_ep_path).module()
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
     cos_sim = cosine_similarity(outputs_pyt, outputs_trt[0])
@@ -362,6 +368,7 @@ def forward(self, x):
     torchtrt.save(trt_module, trt_ep_path)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
 
@@ -420,6 +427,7 @@ def forward(self, x):
     torchtrt.save(trt_module, trt_ep_path)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
+    model.cuda()
     outputs_pyt = model(input)
     outputs_trt = trt_module(input)
 

tests/py/dynamo/runtime/test_004_weight_streaming.py

@@ -69,7 +69,7 @@ def test_weight_streaming_default(self, _, use_python_runtime):
 
         # Weight streaming budget is reverted after the exit from weight streaming context
         assert weight_streaming_ctx.device_budget == requested_budget
-
+        model.cuda()
         ref = model(*input)
         out = optimized_model(*input)
 
@@ -129,7 +129,7 @@ def test_weight_streaming_manual(self, _, use_python_runtime):
             assert weight_streaming_ctx.device_budget == requested_budget
 
             out = optimized_model(*input)
-
+        model.cuda()
         ref = model(*input)
         torch.testing.assert_close(
             out.cpu(),
@@ -168,6 +168,7 @@ def test_weight_streaming_invalid_usage(self, _, use_python_runtime, multi_rt):
         )
 
         # Setting weight streaming context to unsupported module
+        model.cuda()
         with torchtrt.runtime.weight_streaming(model) as weight_streaming_ctx:
             streamable_budget = weight_streaming_ctx.total_device_budget
             assert streamable_budget == 0
@@ -222,7 +223,7 @@ def test_weight_streaming_multi_rt(self, _, use_python_runtime):
                 # Budget distribution to multiple submodule may result in integer differences of at most 1
                 assert abs(weight_streaming_ctx.device_budget - requested_budget) <= 1
                 out = optimized_model(*input)
-
+        model.cuda()
         ref = model(*input)
         torch.testing.assert_close(
             out.cpu(),
@@ -273,7 +274,7 @@ def test_weight_streaming_cudagraphs(self, _, use_python_runtime):
                 weight_streaming_ctx.device_budget = requested_budget
                 for _ in range(4):
                     out = cudagraphs_module(*input)
-
+        model.cuda()
         ref = model(*input)
         torch.testing.assert_close(
             out.cpu(),
@@ -356,7 +357,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
             exp_program,
             **compile_spec,
         )
-
+        model.cuda()
         # List of tuples representing different configurations for three features:
         # Cuda graphs, pre-allocated output buffer, weight streaming change
         states = list(itertools.product((True, False), repeat=3))