feng2022 commited on
Commit
a5397b4
1 Parent(s): 9054fd3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -64,7 +64,7 @@ def load_model(file_name: str, path:str, device: torch.device) -> nn.Module:
64
  f'{file_name}',
65
  use_auth_token=TOKEN)
66
  with open(path, 'rb') as f:
67
- model = pickle.load(f)['G_ema']
68
  model.eval()
69
  model.to(device)
70
  with torch.inference_mode():
 
64
  f'{file_name}',
65
  use_auth_token=TOKEN)
66
  with open(path, 'rb') as f:
67
+ model = torch.load(f)
68
  model.eval()
69
  model.to(device)
70
  with torch.inference_mode():