add test for qkv_size

ZhiyuanChen · ZhiyuanChen · commit 1bbb11ed7b59 · 2022-12-13T20:47:33.000+08:00
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -11664,6 +11664,23 @@ def test_skip_init(self, device):
         self.assertEqual(m_initialized.weight.device, m_uninitialized.weight.device)
         self.assertFalse(torch.allclose(m_initialized.weight, m_uninitialized.weight))
 
+    @dtypes(torch.float)
+    @dtypesIfCUDA(torch.double, torch.float, torch.half)
+    def test_multihead_attention_qkv_diff_size(self, device, dtype):
+        embed_dim = 128
+        k_dim = 64
+        v_dim = 32
+        num_heads = 8
+        sl = 10
+        bs = 8
+        model = nn.MultiheadAttention(embed_dim, num_heads, kdim=k_dim, vdim=v_dim).to(device).to(dtype)
+        q = torch.randn(sl, bs, embed_dim, device=device, dtype=dtype)
+        k = torch.randn(sl, bs, k_dim, device=device, dtype=dtype)
+        v = torch.randn(sl, bs, v_dim, device=device, dtype=dtype)
+        out = model(q, k, v)
+        self.assertEqual(q.size(), out[0].size())
+        self.assertEqual(dtype, out[0].dtype)
+
     @dtypes(torch.float)
     @dtypesIfCUDA(torch.double, torch.float, torch.half)
     def test_transformerencoderlayer(self, device, dtype):
