Update app.py

app.py

@@ -182,31 +182,19 @@ def train_main(train_loader, test_loader, num_epochs, optimizer, model, device='
         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')
+    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:
+    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')
-    axs[0].plot(test_losses, label='Test Loss')
-    axs[0].set_xlabel('Epoch')
-    axs[0].set_ylabel('Loss')
-    axs[0].legend()
-    axs[1].plot(train_accs, label='Train Accuracy')
-    axs[1].plot(test_accs, label='Test Accuracy')
-    axs[1].set_xlabel('Epoch')
-    axs[1].set_ylabel('Accuracy')
-    axs[1].legend()
-    plt.savefig('loss_and_accuracy.png')
+    
     
 
-num_epochs=3
+num_epochs=1
 train_main(train_loader, test_loader, num_epochs, optimizer, model, device)