Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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,
|