Spaces:
Sleeping
Sleeping
Upload folder using huggingface_hub
Browse files
app.py
CHANGED
|
@@ -24,7 +24,9 @@ class_labels= ['Beagle', 'Boxer', 'Bulldog', 'Dachshund', 'German_Shepherd', 'Go
|
|
| 24 |
|
| 25 |
|
| 26 |
# Model
|
| 27 |
-
model = torch.jit.load('best_model.pt').to(device)
|
|
|
|
|
|
|
| 28 |
|
| 29 |
@torch.no_grad()
|
| 30 |
def predict_fn(img:Image):
|
|
|
|
| 24 |
|
| 25 |
|
| 26 |
# Model
|
| 27 |
+
if device=='cuda': model = torch.jit.load('best_model.pt').to(device)
|
| 28 |
+
else: model=torch.load('best_model.pt')
|
| 29 |
+
|
| 30 |
|
| 31 |
@torch.no_grad()
|
| 32 |
def predict_fn(img:Image):
|