Update tsr/system.py

tsr/system.py

@@ -69,7 +69,7 @@ class TSR(BaseModule):
         ckpt = torch.load(weight_path, map_location="cpu")
         
         if "module" in list(ckpt["state_dict"].keys())[0]:
-            state_dict = {key.replace('module.',''): item for key, item in checkpoint["state_dict"].items()}
+            state_dict = {key.replace('module.',''): item for key, item in ckpt["state_dict"].items()}
         else:
             state_dict = ckpt["state_dict"]
         missing_states = set(model.state_dict().keys()) - set(state_dict.keys())