Fix missing return of "logit_bias" in CoCa.forward

mlfoundations · Oct 31, 2023 · 596fbe7 · 596fbe7
1 parent 6f05e1d
commit 596fbe7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/open_clip/coca_model.py b/src/open_clip/coca_model.py
@@ -173,13 +173,16 @@ def forward(
         labels = text[:, -token_embs.shape[1]:]
 
         logits = self.text_decoder(image_embs, token_embs)
-        return {
+        out_dict = {
             "image_features": image_latent,
             "text_features": text_latent,
             "logits": logits,
             "labels": labels,
             "logit_scale": self.logit_scale.exp()
         }
+        if self.logit_bias is not None:
+            out_dict["logit_bias"] = self.logit_bias
+        return out_dict
 
     def generate(
         self,