register OpOverload when OpOverloadPacket is registered

pytorch · Feb 21, 2024 · 2cd4c4c · 2cd4c4c
1 parent 0192ff7
commit 2cd4c4c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/ops_registry.py b/experimental/torch_xla2/torch_xla2/ops_registry.py
@@ -24,6 +24,9 @@ def _lookup(self, op):
     return candidate
 
   def register(self, op, lowering):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      if hasattr(op, 'default'):
+        self.registered_ops[op.default] = lowering
     self.registered_ops[op] = lowering