changing the device setting in conversion.py

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -10,13 +10,14 @@
 from torch_tensorrt._features import ENABLED_FEATURES
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
+from torch_tensorrt.dynamo._defaults import default_device
 from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo.conversion._TRTInterpreter import (
     TRTInterpreter,
     TRTInterpreterResult,
 )
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
-from torch_tensorrt.dynamo.utils import get_model_device, get_torch_inputs
+from torch_tensorrt.dynamo.utils import get_torch_inputs, to_torch_device
 
 import tensorrt as trt
 
@@ -41,7 +42,7 @@ def infer_module_output_dtypes(
     with unset_fake_temporarily():
         # Get the device on which the model exists
         # For large models, this can be done on CPU to save GPU memory allocation for TRT.
-        device = get_model_device(module)
+        device = to_torch_device(default_device())
         torch_inputs = get_torch_inputs(inputs, device)
         if kwarg_inputs is None:
             kwarg_inputs = {}

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -318,7 +318,7 @@ def reduce_operation_with_scatter(
             print("Invalid Operation for Reduce op!!")
 
         operation_rhs = torch.scatter(scatter_tensor, dim, index_tensor, src_tensor)
-        device = to_torch_device(default_device())
+        device = to_torch_device(scatter_tensor.device)
         operation_lhs = operation_lhs.to(device)
         operation_rhs = operation_rhs.to(device)
         return self.func(operation_lhs, operation_rhs)