🐛 Fixes analyze.sh script

2025-03-23 20:46:28 +01:00
parent ee66b08778
commit 929ba716a0
4 changed files with 13 additions and 2 deletions
@@ -11,8 +11,10 @@ parser = argparse.ArgumentParser(
    prog="train_nn"
 )
 parser.add_argument("-g", "--graph", action="store_true", default=False, help="Graph losses")
+parser.add_argument("-s", "--save", default="", help="Graph save location")
 args = parser.parse_args()
 graph = args.graph
+save = args.save


 class NeuralNetwork(nn.Module):
@@ -179,4 +181,7 @@ if graph:

    plt.legend()
    plt.tight_layout()
-    plt.show()
+    if save != "":
+        plt.savefig(save)
+    else:
+        plt.show()