simplifying as per Ross comment

mlfoundations · Apr 16, 2023 · c48111d · c48111d
1 parent 17149cc
commit c48111d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/training/main.py b/src/training/main.py
@@ -250,6 +250,7 @@ def main(args):
         print(f'=> replacing linear layers with {args.use_bnb_linear}')
         linear_replacement_cls = getattr(bnb.nn.triton_based_modules, args.use_bnb_linear)
         replace_linear(model, linear_replacement_cls)
+        model = model.to(device)
 
     random_seed(args.seed, args.rank)
 
@@ -287,8 +288,6 @@ def main(args):
         if args.ddp_static_graph:
             # this doesn't exist in older PyTorch, arg only added if enabled
             ddp_args['static_graph'] = True
-        if args.use_bnb_linear is not None:
-            model = model.to(device)
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device], **ddp_args)
 
         if args.distill: