move import back

Author: Seventeen17

apps/train.py

@@ -1,15 +1,16 @@
 import torch
 
+from flashmodels import Builder, Trainer, accelerate, arguments
+
+
 def train():
     torch.manual_seed(101)
 
     # parse args
-    from flashmodels import arguments
     args = arguments.parse()
 
     # build model, tokenizer, loader, optimizer and lr_scheduler
     # and use accelerator to speed up training
-    from flashmodels import Builder, Trainer, accelerate
     builder = Builder(args)
     model, loader, tokenizer = builder.build_model_dataloader()
     model, loader = accelerate(model, loader, args)