depth and heighte CLI parameters

train.py

@@ -1308,6 +1308,9 @@ if __name__ == "__main__":
     t1_parser.add_argument("--bs", type=int, default=128)
     t1_parser.add_argument("-e", "--epochs", type=int, default=100)
 
+    t1_parser.add_argument("--height", type=int, default=1)
+    t1_parser.add_argument("--depth", type=int, default=1)
+
     t1_parser.add_argument("--attention", type=str, default="self-attention")
     t1_parser.add_argument("--position", type=str, default="fixed")
     t1_parser.add_argument("--dont-normalize", action="store_true")
@@ -1333,6 +1336,9 @@ if __name__ == "__main__":
     t1fs_parser.add_argument("--bs", type=int, default=128)
     t1fs_parser.add_argument("-e", "--epochs", type=int, default=100)
 
+    t1fs_parser.add_argument("--height", type=int, default=1)
+    t1fs_parser.add_argument("--depth", type=int, default=1)
+
     t1fs_parser.add_argument("--attention", type=str, default="self-attention")
     t1fs_parser.add_argument("--position", type=str, default="fixed")
     t1fs_parser.add_argument("--dont-normalize", action="store_true")
@@ -1344,6 +1350,7 @@ if __name__ == "__main__":
     t1fs_parser.add_argument("--epochs-between-save", default=10, type=int)
     t1fs_parser.add_argument("--show-progress", default=50, type=int)
     t1fs_parser.add_argument("--sample", action="store_true")
+    t1fs_parser.set_defaults(freeze_attention=False)
     t1fs_parser.set_defaults(func=t1_from_scratch)
 
     t2_parser = subparsers.add_parser("t2")
@@ -1358,6 +1365,9 @@ if __name__ == "__main__":
     t2_parser.add_argument("--bs", type=int, default=128)
     t2_parser.add_argument("-e", "--epochs", type=int, default=100)
 
+    t2_parser.add_argument("--height", type=int, default=1)
+    t2_parser.add_argument("--depth", type=int, default=1)
+
     t2_parser.add_argument("--attention", type=str, default="self-attention")
     t2_parser.add_argument("--position", type=str, default="fixed")
     t2_parser.add_argument("--dont-normalize", action="store_true")
@@ -1383,6 +1393,9 @@ if __name__ == "__main__":
     t2fs_parser.add_argument("--bs", type=int, default=128)
     t2fs_parser.add_argument("-e", "--epochs", type=int, default=100)
 
+    t2fs_parser.add_argument("--height", type=int, default=1)
+    t2fs_parser.add_argument("--depth", type=int, default=1)
+
     t2fs_parser.add_argument("--attention", type=str, default="self-attention")
     t2fs_parser.add_argument("--position", type=str, default="fixed")
     t2fs_parser.add_argument("--dont-normalize", action="store_true")
@@ -1394,6 +1407,7 @@ if __name__ == "__main__":
     t2fs_parser.add_argument("--epochs-between-save", default=10, type=int)
     t2fs_parser.add_argument("--show-progress", default=50, type=int)
     t2fs_parser.add_argument("--sample", action="store_true")
+    t2fs_parser.set_defaults(freeze_attention=False)
     t2fs_parser.set_defaults(func=t2_from_scratch)
 
     cam_t1_parser = subparsers.add_parser("camembert-t1")