Update visual.py

visual.py

@@ -189,7 +189,7 @@ class VisualAttention(nn.Module):
         # query/key/value: [sq, b, h]
         sq, b, _ = query.size()
 
-        assert query is key, 'Only Support Self-Attention Currently'
+        assert torch.allclose(query, key), 'Only Support Self-Attention Currently'
         sk = sq
         mixed_x_layer = self.in_proj(query)