_refit: Properly compare device type (#3149)

pytorch · Sep 10, 2024 · 91b8f6c · 91b8f6c
1 parent 8154408
commit 91b8f6c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/py/torch_tensorrt/dynamo/_refit.py b/py/torch_tensorrt/dynamo/_refit.py
@@ -156,7 +156,7 @@ def _refit_single_trt_engine_with_gm(
         # Get the refitting mapping
         trt_wt_location = (
             trt.TensorLocation.DEVICE
-            if torch_device == "cuda"
+            if torch_device.type == "cuda"
             else trt.TensorLocation.HOST
         )
         mapping = construct_refit_mapping_from_weight_name_map(