Update.

[picoclvr.git] / main.py

diff --git a/main.py b/main.py
index 00e19ac..704dff5 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -99,6 +99,11 @@ parser.add_argument("--rpl_nb_runs", type=int, default=5)
 
 parser.add_argument("--rpl_no_prog", action="store_true", default=False)
 
+##############################
+# grid options
+
+parser.add_argument("--grid_size", type=int, default=6)
+
 ##############################
 # picoclvr options
 
@@ -517,8 +522,7 @@ elif args.task == "grid":
         nb_train_samples=args.nb_train_samples,
         nb_test_samples=args.nb_test_samples,
         batch_size=args.batch_size,
-        height=args.picoclvr_height,
-        width=args.picoclvr_width,
+        size=args.grid_size,
         logger=log_string,
         device=device,
     )