Update visual.py

visual.py

@@ -125,7 +125,7 @@ class Resampler(nn.Module):
         self.ln_q = norm_layer(embed_dim)
         self.ln_kv = norm_layer(embed_dim)
         
-        self.apply(self._init_weights)
+        # self.apply(self._init_weights)
 
     def _init_weights(self, m):
         if isinstance(m, nn.Linear):
@@ -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)