Spaces:

eternalBlissard
/

FoodVision101

Runtime error

eternalBlissard commited on Oct 21, 2024

Commit

55e8d00

verified ·

1 Parent(s): 71c2fd4

Update predictor.py

Files changed (1) hide show

predictor.py CHANGED Viewed

@@ -18,6 +18,7 @@ def predictionMaker(img):
   ViTModel.to('cuda')
   startTime = timer()
   img = VitTransforms(img).unsqueeze(0)
   ViTModel.eval()
   with torch.inference_mode():
     predProbs = torch.softmax(ViTModel(img),dim=1)

   ViTModel.to('cuda')
   startTime = timer()
   img = VitTransforms(img).unsqueeze(0)
+  img = img.to('cuda')
   ViTModel.eval()
   with torch.inference_mode():
     predProbs = torch.softmax(ViTModel(img),dim=1)