SkalskiP commited on
Commit
2619d65
1 Parent(s): 9364ec8
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -35,7 +35,7 @@ def classify_image(input_image) -> str:
35
  text=imagenet_classes,
36
  images=input_image,
37
  return_tensors="pt",
38
- padding=True)
39
  outputs = model(**inputs)
40
  probs = outputs.logits_per_image.softmax(dim=1)
41
  class_index = np.argmax(probs.detach().numpy())
 
35
  text=imagenet_classes,
36
  images=input_image,
37
  return_tensors="pt",
38
+ padding=True).to(device)
39
  outputs = model(**inputs)
40
  probs = outputs.logits_per_image.softmax(dim=1)
41
  class_index = np.argmax(probs.detach().numpy())