yinuozhang commited on
Commit
3008831
·
verified ·
1 Parent(s): 36f2fbb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -89,9 +89,9 @@ if weight_path is None:
89
  model = MultitaskProteinModel(config)
90
  if weight_path.endswith(".safetensors"):
91
  from safetensors.torch import load_file
92
- state_dict = load_file(weight_path, device="cpu")
93
  else:
94
- state_dict = torch.load(weight_path, map_location="cpu")
95
  missing, unexpected = model.load_state_dict(state_dict, strict=False)
96
  if missing or unexpected:
97
  print(f"[load_state_dict] missing={len(missing)} unexpected={len(unexpected)}")
 
89
  model = MultitaskProteinModel(config)
90
  if weight_path.endswith(".safetensors"):
91
  from safetensors.torch import load_file
92
+ state_dict = load_file(weight_path, device="cpu", weights_only=False)
93
  else:
94
+ state_dict = torch.load(weight_path, map_location="cpu", weights_only=False)
95
  missing, unexpected = model.load_state_dict(state_dict, strict=False)
96
  if missing or unexpected:
97
  print(f"[load_state_dict] missing={len(missing)} unexpected={len(unexpected)}")