X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=cae20f89347b076226b6525d561dab1d89a55e8b;hb=43fbfaac1850098f5b1a9470c8e6ca3d5ab479fe;hp=fabebddcffc31b5c383f2cacde6b8016e03ea1c5;hpb=664758db86b059b68cd11e889a20cc9681e4324a;p=mygptrnn.git

diff --git a/main.py b/main.py
index fabebdd..cae20f8 100755
--- a/main.py
+++ b/main.py
@@ -24,6 +24,17 @@ else:
 
 ######################################################################
 
+
+def str2bool(x):
+    x = x.lower()
+    if x in {"1", "true", "yes"}:
+        return True
+    elif x in {"0", "false", "no"}:
+        return False
+    else:
+        raise ValueError
+
+
 parser = argparse.ArgumentParser(
     description="An implementation of GPT with cache.",
     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -68,7 +79,7 @@ parser.add_argument("--min_learning_rate", type=float, default=6e-5)
 
 # legacy
 
-parser.add_argument("--legacy_lr_schedule", action="store_true", default=False)
+parser.add_argument("--legacy_lr_schedule", type=str2bool, default=True)
 
 parser.add_argument("--legacy_large_lr", type=float, default=1e-4)
 
@@ -478,7 +489,7 @@ def get_lr(n_epoch, it):
 
         if it < args.nb_warmup_iter:
             return args.legacy_large_lr * it / args.nb_warmup_iter
-        elif it < args.legacy_nb_epoch_large_lr:
+        elif n_epoch < args.legacy_nb_epoch_large_lr:
             return args.legacy_large_lr
         else:
             return args.legacy_small_lr