fix:add device

src/streamlit_app.py

@@ -20,6 +20,7 @@ def load_model():
             if re.search(r'in\d+\.running_(mean|var)$', k):
                 del state_dict[k]
         model.load_state_dict(state_dict)
+        model = model.to(device)
         model.eval()
         return model
 
@@ -45,7 +46,7 @@ else:
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         ])
         
-        image_tensor = preprocess(image).unsqueeze(0)
+        image_tensor = preprocess(image).unsqueeze(0).to(device)
         
         with torch.no_grad():
             output = model(image_tensor)