diff --git a/Train_2Photon.py b/Train_2Photon.py
index fa066fd..a8ca505 100644
--- a/Train_2Photon.py
+++ b/Train_2Photon.py
@@ -268,6 +268,6 @@ if __name__ == "__main__":
             print(f"Saved model checkpoint: {modelName}_E{epoch}.pth")
             plot_loss_curves(TrainLosses, ValLosses, test_loss=-1, exp_name=exp_name, conf=conf)
             plot_rms_curve(train_rms_xs, train_rms_ys, val_rms_xs, val_rms_ys, test_rms_x = -1, test_rms_y = -1, exp_name=exp_name, conf=conf)
-    test_loss = evaluate(model, testLoader, loss_fn)
+    test_loss, rms_2d_test = evaluate(model, testLoader, loss_fn)
     plot_loss_curves(TrainLosses, ValLosses, test_loss=test_loss, exp_name=exp_name, conf=conf)
-    plot_rms_curve(train_rms_xs, train_rms_ys, val_rms_xs, val_rms_ys, test_rms_x = -1, test_rms_y = -1, exp_name=exp_name, conf=conf)
\ No newline at end of file
+    plot_rms_curve(train_rms_xs, train_rms_ys, val_rms_xs, val_rms_ys, test_rms_x = rms_2d_test[0], test_rms_y = rms_2d_test[1], exp_name=exp_name, conf=conf)
\ No newline at end of file