update

test/nn/attention/test_performer.py

@@ -5,9 +5,8 @@
 
 def test_performer_attention():
     x = torch.randn(1, 4, 16)
-    mask = torch.ones([1, 4], dtype=torch.bool)
     attn = PerformerAttention(channels=16, heads=4)
-    out = attn(x, mask)
+    out = attn(x)
     assert out.shape == (1, 4, 16)
     assert str(attn) == ('PerformerAttention(heads=4, '
-                         'head_channels=64 kernel=ReLU())')
+                         'head_channels=64)')

test/nn/attention/test_polynormer.py

@@ -3,11 +3,12 @@
 from torch_geometric.nn.attention import PolynormerAttention
 
 
-def test_performer_attention():
+def test_polynormer_attention():
     x = torch.randn(1, 4, 16)
-    mask = torch.ones([1, 4], dtype=torch.bool)
     attn = PolynormerAttention(channels=16, heads=4)
-    out = attn(x, mask)
+    out = attn(x)
+    import pdb
+    pdb.set_trace()
     assert out.shape == (1, 4, 16)
     assert str(attn) == ('PolynormerAttention(heads=4, '
                          'head_channels=64 kernel=ReLU())')

torch_geometric/nn/attention/polynormer.py

@@ -87,5 +87,4 @@ def reset_parameters(self):
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}('
                 f'heads={self.heads}, '
-                f'head_channels={self.head_channels} '
-                f'kernel={self.kernel})')
+                f'head_channels={self.head_channels})')