update

pyg-team · Jan 1, 2025 · ab3b2d9 · ab3b2d9
1 parent e43267c
commit ab3b2d9
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/test/nn/attention/test_polynormer.py b/test/nn/attention/test_polynormer.py
@@ -4,11 +4,9 @@
 
 
 def test_polynormer_attention():
-    x = torch.randn(1, 4, 16)
+    x = torch.randn(4, 16)
     attn = PolynormerAttention(channels=16, heads=4)
     out = attn(x)
-    import pdb
-    pdb.set_trace()
-    assert out.shape == (1, 4, 16)
+    assert out.shape == (4, 256)
     assert str(attn) == ('PolynormerAttention(heads=4, '
-                         'head_channels=64 kernel=ReLU())')
+                         'head_channels=64)')