Convert JIT model (on state dict load) to sd for pretrained='filename…

….pt' support for OpenAI .pt files. Fix #622
mlfoundations · Sep 12, 2023 · a4b5bc3 · a4b5bc3
1 parent 6be6dc8
commit a4b5bc3
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
@@ -88,6 +88,10 @@ def load_state_dict(checkpoint_path: str, map_location='cpu'):
     checkpoint = torch.load(checkpoint_path, map_location=map_location)
     if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
         state_dict = checkpoint['state_dict']
+    elif isinstance(checkpoint, torch.jit.ScriptModule):
+        state_dict = checkpoint.state_dict()
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            state_dict.pop(key, None)
     else:
         state_dict = checkpoint
     if next(iter(state_dict.items()))[0].startswith('module'):