Update __torch_function__ idiom (#599)

Co-authored-by: Fritz Obermeyer <[email protected]>
pyro-ppl · Oct 28, 2022 · 0a8e4b2 · 0a8e4b2
1 parent 1d07af1
commit 0a8e4b2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/funsor/torch/provenance.py b/funsor/torch/provenance.py
@@ -27,7 +27,8 @@ def __init__(self, data, provenance=frozenset()):
     def __repr__(self):
         return "Provenance:\n{}\nTensor:\n{}".format(self._provenance, self._t)
 
-    def __torch_function__(self, func, types, args=(), kwargs=None):
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
         if kwargs is None:
             kwargs = {}
         # collect provenance information from args