Update app.py

app.py

@@ -82,9 +82,8 @@ class DummyModel(nn.Module):
         return dec4
 
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = DummyModel()
-model.load_state_dict(torch.load('model.pth'))
+model_weights = torch.load('model.pth', map_location=device)
+model.load_state_dict(model_weights)
 model = model.to(device)
 model.eval()