Update app.py

app.py

@@ -181,12 +181,16 @@ def train_main(train_loader, test_loader, num_epochs, optimizer, model, device='
         train_accs.append(train_acc)
         test_accs.append(test_acc)
         print('Epoch: {}, train loss: {:.4f}, test loss: {:.4f}, train accuracy: {:.2f}%, test accuracy: {:.2f}%'.format(epoch+1, train_loss, test_loss, train_acc, test_acc))
-        
+    print("Saving the model")
+    torch.save(model, 'trained_model.pt')
+    print("Model saved successfully")
         # save the results to a text file
+
+    print("Saving training logs")
     with open("results.txt", "w") as f:
         for epoch in range(num_epochs):
             f.write("Epoch: {}, train loss: {:.4f}, test loss: {:.4f}, train accuracy: {:.2f}%, test accuracy: {:.2f}%\n".format(epoch+1, train_losses[epoch], test_losses[epoch], train_accs[epoch], test_accs[epoch]))
-
+    print("Logs saved")
     # plot the loss and accuracy curves side by side
     fig, axs = plt.subplots(1, 2, figsize=(12, 4))
     axs[0].plot(train_losses, label='Train Loss')
@@ -200,8 +204,7 @@ def train_main(train_loader, test_loader, num_epochs, optimizer, model, device='
     axs[1].set_ylabel('Accuracy')
     axs[1].legend()
     plt.savefig('loss_and_accuracy.png')
-    torch.save(model, 'trained_model.pt')
-    plt.show()
+    
 
 num_epochs=3
 train_main(train_loader, test_loader, num_epochs, optimizer, model, device)