clearer mae

lucidrains · Nov 22, 2021 · 9f8c606 · 9f8c606
1 parent 5ae5557
commit 9f8c606
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vit_pytorch/mae.py b/vit_pytorch/mae.py
@@ -78,12 +78,12 @@ def forward(self, img):
 
         # concat the masked tokens to the decoder tokens and attend with decoder
 
-        decoder_tokens = torch.cat((decoder_tokens, mask_tokens), dim = 1)
+        decoder_tokens = torch.cat((mask_tokens, decoder_tokens), dim = 1)
         decoded_tokens = self.decoder(decoder_tokens)
 
         # splice out the mask tokens and project to pixel values
 
-        mask_tokens = decoded_tokens[:, -num_masked:]
+        mask_tokens = decoded_tokens[:, :num_masked]
         pred_pixel_values = self.to_pixels(mask_tokens)
 
         # calculate reconstruction loss