Spaces:

erow
/

SimilarityThresholding

Sleeping

Gent (PG/R - Comp Sci & Elec Eng) commited on Aug 12, 2023

Commit

fe6bd89

1 Parent(s): bc8c24d

load weights

Files changed (1) hide show

utils.py CHANGED Viewed

@@ -48,7 +48,8 @@ def load_pretrained_weights(model, pretrained_weights, checkpoint_key=None, pref
     print("Load pre-trained checkpoint from: %s[%s] at %d epoch" % (pretrained_weights, checkpoint_key, epoch))
-    state_dict = state_dict[checkpoint_key]
     # remove `module.` prefix
     if prefixes is None: prefixes= ["module.","backbone."]
     for prefix in prefixes:

     print("Load pre-trained checkpoint from: %s[%s] at %d epoch" % (pretrained_weights, checkpoint_key, epoch))
+    if checkpoint_key:
+        state_dict = state_dict[checkpoint_key]
     # remove `module.` prefix
     if prefixes is None: prefixes= ["module.","backbone."]
     for prefix in prefixes: