make as with cls at the end

mlfoundations · Jun 24, 2023 · 852b2f4 · 852b2f4
1 parent fb72f4d
commit 852b2f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/open_clip/transformer.py b/src/open_clip/transformer.py
@@ -499,7 +499,7 @@ def forward(self, x: torch.Tensor):
 
         if self.output_tokens:
             return pooled, tokens
-        
+
         return pooled
 
 
@@ -586,7 +586,7 @@ def build_attention_mask(self):
 
     def build_cls_mask(self, text, cast_dtype: torch.dtype):
         cls_mask = (text != self.pad_id).unsqueeze(1)
-        cls_mask = F.pad(cls_mask, (1, 0, cls_mask.shape[2], 0), value=1.0)
+        cls_mask = F.pad(cls_mask, (0, 1, cls_mask.shape[2], 0), value=1.0)
         additive_mask = torch.empty(cls_mask.shape, dtype=cast_dtype, device=cls_mask.device)
         additive_mask.fill_(0)
         additive_mask.masked_fill_(~cls_mask, float("-inf"))