cbensimon HF staff commited on
Commit
1433399
1 Parent(s): c586b21

Enable GPU when available

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -10,7 +10,7 @@ torch.set_grad_enabled(False)
10
 
11
  from model import Generator
12
 
13
- device = "cpu"
14
 
15
  model = Generator().eval().to(device)
16
  model.load_state_dict(torch.load("face_paint_512_v2_0.pt"))
 
10
 
11
  from model import Generator
12
 
13
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
 
15
  model = Generator().eval().to(device)
16
  model.load_state_dict(torch.load("face_paint_512_v2_0.pt"))