Update tsr/system.py

tsr/system.py

@@ -67,7 +67,16 @@ class TSR(BaseModule):
         OmegaConf.resolve(cfg)
         model = cls(cfg)
         ckpt = torch.load(weight_path, map_location="cpu")
-        model.load_state_dict(ckpt)
+        
+        if "module" in list(ckpt["state_dict"].keys())[0]:
+            state_dict = {key.replace('module.',''): item for key, item in checkpoint["state_dict"].items()}
+        else:
+            state_dict = ckpt["state_dict"]
+        missing_states = set(model.state_dict().keys()) - set(state_dict.keys())
+        if len(missing_states) > 0:
+            warnings.warn("Missing keys ! : {}".format(missing_states))
+            
+        model.load_state_dict(state_dict)
         return model
 
     def configure(self):