Added warning and changed the test cases

Author: cehongwang

py/torch_tensorrt/dynamo/backend/backends.py

@@ -10,7 +10,6 @@
 from torch._dynamo.backends.common import aot_autograd
 from torch._dynamo.utils import detect_fake_mode
 from torch._functorch.aot_autograd import aot_export_joint_simple
-from torch.distributed.tensor import DTensor
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo._compiler import compile_module
 from torch_tensorrt.dynamo.lowering import (
@@ -89,6 +88,11 @@ def aot_torch_tensorrt_aten_backend(
             logger.warning(
                 "It is recommended to run the model with use_distributed_mode_trace = True since there are distributed tensors in the input which is not supported in aot_export_joint_simple"
             )
+
+    if settings.offload_module_to_cpu:
+        logger.warning(
+            "`offload_module_to_cpu` is not supported for `torch_compile` backend."
+        )
     return _pretraced_backend(gm, sample_inputs, settings, engine_cache)
 
 

tests/py/dynamo/models/test_models.py

@@ -79,11 +79,12 @@ def test_resnet18_cpu_offload(ir):
     }
 
     trt_mod = torchtrt.compile(model, **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()
     cos_sim = cosine_similarity(model(input), trt_mod(input))
     assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
@@ -287,11 +288,12 @@ def test_bert_base_uncased_cpu_offload(ir):
         "offload_module_to_cpu": True,
     }
     trt_mod = torchtrt.compile(model, **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()
 
     model_outputs = model(input, input2)
     trt_model_outputs = trt_mod(input, input2)