test fix

pytorch · Sep 12, 2024 · d8fb3de · d8fb3de
1 parent 7762dd4
commit d8fb3de
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -391,6 +391,7 @@ def setUpClass(cls):
 
   def setUp(self):
     self.env = tensor.Environment()
+    torch.manual_seed(0)
 
   # Replaces all values in the input torch_tensor that are less than the given threshold
   # with the threshold value itself.

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -57,6 +57,7 @@ def op(*aten, **kwargs):
   def inner(func):
     for a in aten:
       ops_registry.register_torch_dispatch_op(a, func, **kwargs)
+      continue
 
       if isinstance(a, torch._ops.OpOverloadPacket):
         opname = a.default.name() if 'default' in a.overloads() else a._qualified_op_name
@@ -437,6 +438,14 @@ def _aten_dot(x, y):
   return jnp.dot(x, y)
 
 
+@op(torch.ops.aten._to_copy)
+def _aten__to_copy(self, **kwargs):
+  dtype = mappings.t2j_dtype(kwargs["dtype"])
+  if dtype != self.dtype:
+    return self.astype(dtype)
+  return jnp.copy(self)
+
+
 
 
 @op(torch.ops.aten.empty)