Update app.py

app.py

@@ -80,6 +80,9 @@ for epoch in range(num_epochs):
     print(f'Validation Accuracy: {100 * correct / total:.2f}%')
 
 import wandb
+torch.save(model.state_dict(), 'weather_model.pth')
+wandb.save('weather_model.pth')
+
 def predict_image(image):
     image = Image.fromarray(image).convert("RGB")
     image = transform(image)
@@ -112,5 +115,3 @@ interface = gr.Interface(fn=predict_image,
 interface.launch()
 
 # Save model checkpoint to W&B
-torch.save(model.state_dict(), 'weather_model.pth')
-wandb.save('weather_model.pth')