diff options
Diffstat (limited to 'ep_run/casc_bp_train.py')
| -rw-r--r-- | ep_run/casc_bp_train.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index d45327b..ebd7536 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -25,12 +25,13 @@ ap.add_argument('--resume', default='') # path to a ckpt (tok/p ap.add_argument('--olmo2', action='store_true') # OLMo2-standard block (see casc_eq_train.py) ap.add_argument('--wd', type=float, default=-1.0) # >=0: grouped weight decay; <0 = legacy uniform 1e-4 ap.add_argument('--zloss', type=float, default=0.0) # z-loss coefficient; 0 = off +ap.add_argument('--data', default='tinystories_bpe') # dataset dir under ep_run/data args = ap.parse_args() if args.olmo2 and args.tok_init <= 0: args.tok_init = 0.02 torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' -DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe') +DD = Path('/home/yurenh2/ept/ep_run/data') / args.data vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size'] def get_batch(split): |