_refit: Properly compare device type (#3149)

HolyWu · web-flow · commit 91b8f6cf02b3 · 2024-09-10T13:58:20.000-07:00
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(

Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ def _refit_single_trt_engine_with_gm(`
`156`	`156`	`# Get the refitting mapping`
`157`	`157`	`trt_wt_location = (`
`158`	`158`	`trt.TensorLocation.DEVICE`
`159`		`- if torch_device == "cuda"`
	`159`	`+ if torch_device.type == "cuda"`
`160`	`160`	`else trt.TensorLocation.HOST`
`161`	`161`	`)`
`162`	`162`	`mapping = construct_refit_mapping_from_weight_name_map(`