Update jaten.py

pytorch · Jul 18, 2024 · f64633e · f64633e
1 parent 472d194
commit f64633e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -284,6 +284,8 @@ def _aten_div(x, y, rounding_mode=""):
 
   if rounding_mode == "floor":
     res = jnp.floor_divide(x, y)
+    if _is_int(x) and _is_int(y):
+      res_dtype = jnp.dtype('int64')
   else:
     res = x / y
   if rounding_mode == "trunc":