test: add test for attentional pooling layer

tests/models/coca/test_coca.py

@@ -3,7 +3,7 @@
 import torch
 
 from modalities.__main__ import load_app_config_dict
-from modalities.models.coca.coca_model import CoCa, CoCaConfig
+from modalities.models.coca.coca_model import AttentionalPooling, CoCa, CoCaConfig
 
 _ROOT_DIR = Path(__file__).parents[1]
 
@@ -25,3 +25,11 @@ def test_coca_forward():
     assert out["logits"].shape == (1, 1024, 50304)
     assert out["vision_cls"].shape == (1, 1, 768)
     assert out["text_cls"].shape == (1, 1, 768)
+
+
+def test_attn_pool():
+    model = AttentionalPooling(n_embd=768, n_head=8, bias=False, epsilon=1e-5)
+    dummy_vision_embed = torch.randn(1, 256, 768)
+    dummy_queries = torch.randn(1, 257, 768)
+    out = model(dummy_vision_embed, dummy_queries)
+    assert out.shape == (1, 257, 768)