AxionLab-official commited on
Commit
76df893
·
verified ·
1 Parent(s): 2a7b0a6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -160,7 +160,7 @@ model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
160
 
161
  cfg = Config()
162
  model = StudentUNet(cfg)
163
- ckpt = torch.load(model_path, map_location="cpu")
164
 
165
  # Trata se você salvou apenas o state_dict ou o dicionário inteiro de treino
166
  if "model_state" in ckpt:
 
160
 
161
  cfg = Config()
162
  model = StudentUNet(cfg)
163
+ ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
164
 
165
  # Trata se você salvou apenas o state_dict ou o dicionário inteiro de treino
166
  if "model_state" in ckpt: