Update start.py

janelia-cellmap · Mar 20, 2024 · 8bb5d50 · 8bb5d50
1 parent c128020
commit 8bb5d50
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/dacapo/experiments/starts/start.py b/dacapo/experiments/starts/start.py
@@ -30,7 +30,11 @@ def _set_weights(model, weights, run, criterion, old_head=None, new_head=None):
                     head_weights[key] = weights.model[key]
                 for key in head_keys:
                     weights.model.pop(key, None)
-                model.load_state_dict(weights.model, strict=False)
+                try:
+                    model.load_state_dict(weights.model, strict=True)
+                except:
+                    logger.warning("Unable to load model in strict mode. Loading flexibly.")
+                    model.load_state_dict(weights.model, strict=False)
                 model = match_heads(model, head_weights, old_head, new_head)
             except RuntimeError as e:
                 logger.error(f"ERROR starter matching head: {e}")