SerdarHelli commited on
Commit
fafc0c2
1 Parent(s): 8721178

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -38,7 +38,7 @@ network_pkl='braingan-400.pkl'
38
  with open(network_pkl, 'rb') as f:
39
  G = pickle.load(f)['G_ema']
40
 
41
- device = torch.device('cuda')
42
  G.eval()
43
  G.to(device)
44
  def predict(Seed,choices):
 
38
  with open(network_pkl, 'rb') as f:
39
  G = pickle.load(f)['G_ema']
40
 
41
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
42
  G.eval()
43
  G.to(device)
44
  def predict(Seed,choices):