Update codebase/inference/inference.py

codebase/inference/inference.py

@@ -24,12 +24,12 @@ def build_model(model_name, ckpt_path, device):
     if model_name == "ViT-B-32":
         model, _, _ = open_clip.create_model_and_transforms("ViT-B/32", pretrained="openai")
         checkpoint = torch.load(ckpt_path, map_location="cpu")
-        msg = model.load_state_dict(checkpoint, strict=False)
+        msg = model.load_state_dict(checkpoint)
 
     elif model_name == "ViT-H-14":
         model, _, _ = open_clip.create_model_and_transforms("ViT-H/14", pretrained="laion2b_s32b_b79k")
         checkpoint = torch.load(ckpt_path, map_location="cpu")
-        msg = model.load_state_dict(checkpoint, strict=False)
+        msg = model.load_state_dict(checkpoint)
 
     print(msg)
     model = model.to(device)