Update jaten.py

pytorch · Sep 25, 2024 · d0d3298 · d0d3298
1 parent 0a256a1
commit d0d3298
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -2324,6 +2324,9 @@ def _aten_neg(x):
 # aten.nonzero
 @op(torch.ops.aten.nonzero)
 def _aten_nonzero(x):
+  if jnp.ndim(x) == 0: # when x is scalar, return torch.tensor([], size=(1, 0), dtype=torch.int64)
+    res = torch.empty(1, 0, dtype=torch.int64)
+    return jnp.array(res.numpy())
   index_tuple = jnp.nonzero(x)
   index_tuple = [jnp.expand_dims(p, -1) for p in index_tuple]
   return jnp.concatenate(index_tuple, axis=-1)