ericup commited on
Commit
fb0a5c2
1 Parent(s): 9177536

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -0
app.py CHANGED
@@ -21,6 +21,13 @@ def predict(filename, model=None, device=None, reduce_labels=True):
21
 
22
  global default_model
23
  assert isinstance(filename, str)
 
 
 
 
 
 
 
24
  print(dict(
25
  filename=filename,
26
  model=model,
 
21
 
22
  global default_model
23
  assert isinstance(filename, str)
24
+
25
+ if device is None:
26
+ if torch.cuda.device_count():
27
+ device = 'cuda'
28
+ else:
29
+ device = 'cpu'
30
+
31
  print(dict(
32
  filename=filename,
33
  model=model,