From 80c8cd29af3d7b6f519e54128fde6bf27ccee0ee Mon Sep 17 00:00:00 2001 From: Ng Chan Way Date: Sun, 26 Nov 2023 13:02:57 +0800 Subject: [PATCH] feat: Allow flexibility in load_state_dict To allow loading state_dict with additional parameters --- openchem/models/openchem_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openchem/models/openchem_model.py b/openchem/models/openchem_model.py index e2fd376..102b55b 100644 --- a/openchem/models/openchem_model.py +++ b/openchem/models/openchem_model.py @@ -83,7 +83,7 @@ def cast_inputs(sample, task, use_cuda): def load_model(self, path): weights = torch.load(path) - self.load_state_dict(weights) + self.load_state_dict(weights, strict=False) def save_model(self, path): torch.save(self.state_dict(), path)