Update models_milan.py

zejiangh · Aug 16, 2022 · 677c981 · 677c981
1 parent 08cc056
commit 677c981
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models_milan.py b/models_milan.py
@@ -119,7 +119,7 @@ def __init__(self, img_size=224, patch_size=16, in_chans=3,
         if use_clip is False:
             self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True) # decoder to patch
         else:
-            self.decoder_pred = nn.Linear(decoder_embed_dim, 768, bias=True) # clip output [batch x seq x 512]
+            self.decoder_pred = nn.Linear(decoder_embed_dim, 512, bias=True) # clip output [batch x seq x 512]
         # --------------------------------------------------------------------------
 
         self.norm_pix_loss = norm_pix_loss