changed the message and added other cpu offload test

Author: cehongwang

py/torch_tensorrt/dynamo/backend/backends.py

@@ -91,7 +91,7 @@ def aot_torch_tensorrt_aten_backend(
 
     if settings.offload_module_to_cpu:
         logger.warning(
-            "`offload_module_to_cpu` is not supported for `torch_compile` backend."
+            "The offload_module_to_cpu option is set, but it is being ignored since the torch_compile backend does not support this feature"
         )
     return _pretraced_backend(gm, sample_inputs, settings, engine_cache)
 

tests/py/dynamo/models/test_export_serde.py

@@ -320,11 +320,12 @@ def test_resnet18_cpu_offload(ir):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    assertions.assertTrue(
-        get_model_device(model).type == "cpu",
-        msg="Model should be offloaded to CPU",
-    )
-    model.cuda()
+    if ir == "dynamo":
+        assertions.assertTrue(
+            get_model_device(model).type == "cpu",
+            msg="Model should be offloaded to CPU",
+        )
+        model.cuda()
     torchtrt.save(trt_module, trt_ep_path)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()