Update splat_attention.py

splat_attention.py

@@ -216,6 +216,7 @@ class SplatAttention(nn.Module):
         return self.drop(self.w_o(out))
 
 
+
 def _demo():
     torch.manual_seed(0)
     dev = "cuda" if torch.cuda.is_available() else "cpu"
@@ -231,7 +232,7 @@ def _demo():
     if dev == "cuda":
         import time
         mha = nn.MultiheadAttention(256, 8, batch_first=True).to(dev)
-        for L, B in [(128, 32), (2048, 2)]:
+        for L, B in [(128, 256), (256, 32), (512, 32), (768, 32), (1024, 32), (1280, 32), (2048, 32), (4096, 32), (8192, 32)]:
             xx = torch.randn(B, L, 256, device=dev, requires_grad=True)
             def t(fn, n=10):
                 for _ in range(3):
@@ -250,5 +251,6 @@ def _demo():
           "head_gates=True for learnable attenuation)")
 
 
+
 if __name__ == "__main__":
     _demo()