bluelike commited on
Commit
d23bea4
1 Parent(s): 616d8e0

Update visual.py

Browse files
Files changed (1) hide show
  1. visual.py +1 -1
visual.py CHANGED
@@ -189,7 +189,7 @@ class VisualAttention(nn.Module):
189
  # query/key/value: [sq, b, h]
190
  sq, b, _ = query.size()
191
 
192
- assert query is key, 'Only Support Self-Attention Currently'
193
  sk = sq
194
  mixed_x_layer = self.in_proj(query)
195
 
 
189
  # query/key/value: [sq, b, h]
190
  sq, b, _ = query.size()
191
 
192
+ assert torch.allclose(query, key), 'Only Support Self-Attention Currently'
193
  sk = sq
194
  mixed_x_layer = self.in_proj(query)
195