Traly commited on
Commit
7eba3dd
β€’
1 Parent(s): a089d4c
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -52,7 +52,7 @@ def load_checkpoint(model, ckpt_path):
52
  new_state_dict[k] = v
53
 
54
  model.load_state_dict(new_state_dict)
55
- # model.cuda()
56
  del checkpoint
57
  torch.cuda.empty_cache()
58
 
@@ -77,8 +77,8 @@ def image_infer(img_PIL):
77
  trainer.model.eval()
78
  img_lr, img_lr_up = get_img_data(img_PIL, hparams, sr_scale=4)
79
 
80
- # img_lr = img_lr.to('cuda')
81
- # img_lr_up = img_lr_up.to('cuda')
82
 
83
  img_sr, _ = trainer.model.sample(img_lr, img_lr_up, img_lr_up.shape)
84
 
 
52
  new_state_dict[k] = v
53
 
54
  model.load_state_dict(new_state_dict)
55
+ model.cuda()
56
  del checkpoint
57
  torch.cuda.empty_cache()
58
 
 
77
  trainer.model.eval()
78
  img_lr, img_lr_up = get_img_data(img_PIL, hparams, sr_scale=4)
79
 
80
+ img_lr = img_lr.to('cuda')
81
+ img_lr_up = img_lr_up.to('cuda')
82
 
83
  img_sr, _ = trainer.model.sample(img_lr, img_lr_up, img_lr_up.shape)
84