wip

haunt98 · Mar 24, 2024 · bade6cb · bade6cb
1 parent 1e69983
commit bade6cb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/gill/models.py b/gill/models.py
@@ -554,10 +554,10 @@ def __init__(self, tokenizer, model_args: Optional[GILLArgs] = None,
       print('Loading decision model...')
       self.decision_model = nn.Sequential(*[
           nn.Dropout(0.5),
-          nn.Linear(4096, 2),
+          nn.Linear(768, 2),
       ])
       mlp_checkpoint = torch.load(decision_model_path)
-      self.decision_model.load_state_dict(mlp_checkpoint['state_dict'], strict=True)
+      self.decision_model.load_state_dict(mlp_checkpoint['state_dict'])
       self.decision_model.eval()
 
   def __call__(self, images: Tensor, tgt_tokens: Optional[Tensor] = None, caption_len: Optional[Tensor] = None,