Fix unbind op.

The test passed with the previous implementation but I created a
slightly more efficient one anyway.

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -2399,10 +2399,7 @@ def _aten_trunc(a):
 @op(torch.ops.aten.unbind)
 @op(torch.ops.aten.unbind_copy)
 def _aten_unbind(a, dim=0):
-  return tuple(
-    _aten_squeeze_dim(jax.lax.index_in_dim(a, i, axis=dim), dim)
-    for i in range(a.shape[dim])
-  )
+  return [lax.index_in_dim(a, i, axis, keepdims=False) for i in range(a.shape[axis])]
 
 
 # NOTE: skip aten.upsample_nearest2d and aten.upsample_bilinear2d