Remove warnings

Author: JeanKossaifi

tltorch/factorized_tensors/core.py

@@ -6,6 +6,7 @@
 from torch import nn
 import numpy as np
 
+import torch
 
 # Author: Jean Kossaifi
 # License: BSD 3 clause
@@ -368,12 +369,14 @@ def normal_(self, mean=0, std=1):
 
     def __repr__(self):
         return f'{self.__class__.__name__}(shape={self.shape}, rank={self.rank})'
-
-    def __torch_function__(self, func, types, args=(), kwargs=None):
+    
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
         if kwargs is None:
             kwargs = {}
 
         args = [t.to_tensor() if hasattr(t, 'to_tensor') else t for t in args]
+        # return super().__torch_function__(func, types, args, kwargs)
         return func(*args, **kwargs)
 
     @property
@@ -549,7 +552,8 @@ def __repr__(self):
         msg += f'rank={self.rank})'
         return msg
 
-    def __torch_function__(self, func, types, args=(), kwargs=None):
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
         if kwargs is None:
             kwargs = {}
 

tltorch/factorized_tensors/tests/test_factorizations.py

@@ -39,7 +39,7 @@ def test_FactorizedTensor(factorization):
         res = fact_tensor[idx]
         if not torch.is_tensor(res):
             res = res.to_tensor()
-        testing.assert_allclose(reconstruction[idx], res)
+        testing.assert_close(reconstruction[idx], res)
 
 
 @pytest.mark.parametrize('factorization', ['BlockTT', 'CP']) #['CP', 'Tucker', 'BlockTT'])
@@ -85,7 +85,7 @@ def test_TensorizedMatrix(factorization, batch_size):
         res = fact_tensor[idx]
         if not torch.is_tensor(res):
             res = res.to_matrix()
-        testing.assert_allclose(reconstruction[idx], res)
+        testing.assert_close(reconstruction[idx], res)
 
 
 @pytest.mark.parametrize('factorization', ['CP', 'TT'])
@@ -104,7 +104,7 @@ def test_transduction(factorization):
 
         indices = [slice(None)]*mode
         for i in range(new_dim):
-            testing.assert_allclose(original_rec, rec[tuple(indices + [i])])
+            testing.assert_close(original_rec, rec[tuple(indices + [i])])
 
 @pytest.mark.parametrize('unsqueezed_init', ['average', 1.2])
 def test_tucker_init_unsqueezed_modes(unsqueezed_init):
@@ -122,7 +122,7 @@ def test_tucker_init_unsqueezed_modes(unsqueezed_init):
         coef = unsqueezed_init
 
     for i in range(4):
-        testing.assert_allclose(rec[:, i], mat*coef)
+        testing.assert_close(rec[:, i], mat*coef)
 
 
 @pytest.mark.parametrize('factorization', ['ComplexCP', 'ComplexTucker', 'ComplexTT', 'ComplexDense'])