davidserra9 commited on
Commit
c7626cc
·
verified ·
1 Parent(s): cd4d5ac

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -57,7 +57,7 @@ config = OmegaConf.load(CONFIG)
57
  device = "cuda" if torch.cuda.is_available() else "cpu"
58
 
59
  config['train']['cuda_visible_device'] = device
60
- model = NamedCurves(config.model).to(device)
61
  model.load_state_dict(torch.load(model_pt)["model_state_dict"])
62
 
63
  def load_img(filename, norm=True,):
 
57
  device = "cuda" if torch.cuda.is_available() else "cpu"
58
 
59
  config['train']['cuda_visible_device'] = device
60
+ model = NamedCurves(config.model, device=device).to(device)
61
  model.load_state_dict(torch.load(model_pt)["model_state_dict"])
62
 
63
  def load_img(filename, norm=True,):