changing the device setting in conversion.py

Author: apbose

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)

py/torch_tensorrt/dynamo/utils.py

@@ -12,6 +12,7 @@
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo import _defaults
+from torch_tensorrt.dynamo._defaults import default_device
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
@@ -190,7 +191,7 @@ def get_model_device(module: torch.fx.GraphModule) -> torch.device:
             break
 
     if device is None:
-        device = torch.device("cpu")
+        device = to_torch_device(default_device())
         logger.warning(
             "Could not detect the device on which the model exists. Assuming the model is on CPU"
         )