To support zero zero-gpu

app.py

@@ -5,9 +5,11 @@ from PIL import Image
 from timm.data import create_transform
 
 
+device = "cuda"
+
 # Prepare the model.
 import models
-model = models.mambaout_femto(pretrained=True) # can change different model name
+model = models.mambaout_femto(pretrained=True).to(device=device) # can change different model name
 model.eval()
 
 # Prepare the transform.
@@ -17,9 +19,9 @@ transform = create_transform(input_size=224, crop_pct=model.default_cfg['crop_pc
 response = requests.get("https://git.io/JJkYN")
 labels = response.text.split("\n")
 
+[email protected]
 def predict(inp):
-  inp = transform(inp).unsqueeze(0)
-
+  inp = transform(inp).unsqueeze(0).to(device=device)
   with torch.no_grad():
     prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
     confidences = {labels[i]: float(prediction[i]) for i in range(1000)}