--epochs-between-save

train.py

@@ -21,7 +21,7 @@ if __name__ == "__main__":
     parser.add_argument("--device", type=str, default="cpu")
     parser.add_argument("--logdir", type=str, required=False)
     parser.add_argument("-c", "--checkpoint", type=str, required=False)
-    parser.add_argument("--init-w2v", action="store_true")
+    parser.add_argument("--epochs-between-save", default=10, type=int)
 
     position_mapper = {
         "none": PositionalEmbeddingType.NONE,
@@ -127,7 +127,7 @@ if __name__ == "__main__":
         writer.add_scalar("Loss/train", cumloss / len(train_loader), epoch)
         print(f"EPOCH {epoch:04} - Loss: {cumloss / len(train_loader)}")
 
-        if epoch % 10 == 0:
+        if epoch % args.epochs_between_save == 0:
             model.eval()
             k = 5
             train_precision, train_tp = eval_model_k(