+def get_lr(n_epoch, it):
+ if args.legacy_lr_schedule:
+ # my crude scheduling to compare to previous baseline, added
+ # warmup though
+
+ if it < args.nb_warmup_iter:
+ return args.legacy_large_lr * it / args.nb_warmup_iter
+ elif n_epoch < args.legacy_nb_epoch_large_lr:
+ return args.legacy_large_lr
+ else:
+ return args.legacy_small_lr
+
+ # from nanoGPT