update model saving logic

Author: dohlee

.gitignore

@@ -1,6 +1,7 @@
 spliceai_train_code/
 
 data/
+ckpts/
 note/.ipynb_checkpoints/
 
 *.egg-info/

tranception_pytorch/train.py

@@ -136,6 +136,9 @@ def main():
             running_loss = []
 
         if (cnt // args.gradient_accumulation_steps) % 25000 == 0:
+            if not os.path.exists(os.path.dirname(args.output)):
+                os.makedirs(os.path.dirname(args.output))
+
             idx = cnt // args.gradient_accumulation_steps
             torch.save(model.state_dict(), f'{args.output}_{idx}.pt')