Dont wrap ints in torch.tensor; (#6539)

pytorch · Feb 15, 2024 · 911485e · 911485e
1 parent 29b01c6
commit 911485e
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torch_xla/stablehlo.py b/torch_xla/stablehlo.py
@@ -248,9 +248,8 @@ def call_function(self, target, args: Tuple, kwargs: Dict) -> Any:
         args = tuple(
             map(
                 lambda arg_spec: torch.tensor(arg_spec[0])
-                if isinstance(arg_spec[0], (float, int)) and type(arg_spec[
-                    1].type) == torch.TensorType else arg_spec[0],
-                args_and_specs))
+                if isinstance(arg_spec[0], float) and type(arg_spec[1].type) ==
+                torch.TensorType else arg_spec[0], args_and_specs))
     return super().call_function(target, args, new_kwargs)
 
   def run_node(self, n) -> Any: