Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -29,11 +29,11 @@ def inference(input_image):
|
|
| 29 |
#transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| 30 |
])
|
| 31 |
input_tensor = preprocess(input_image)
|
| 32 |
-
|
| 33 |
|
| 34 |
# move the input and model to GPU for speed if available
|
| 35 |
if torch.cuda.is_available():
|
| 36 |
-
|
| 37 |
model.to('cuda')
|
| 38 |
|
| 39 |
with torch.no_grad():
|
|
|
|
| 29 |
#transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| 30 |
])
|
| 31 |
input_tensor = preprocess(input_image)
|
| 32 |
+
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
|
| 33 |
|
| 34 |
# move the input and model to GPU for speed if available
|
| 35 |
if torch.cuda.is_available():
|
| 36 |
+
input_batch = input_batch.to('cuda')
|
| 37 |
model.to('cuda')
|
| 38 |
|
| 39 |
with torch.no_grad():
|