Solved by Cade_Gordon2 in post #2 
              
              
                The trick is parameter groups! 
optimizer = ZeroRedundancyOptimizer(
                rest_params,
                optim.AdamW,
                lr=args.lr,
                betas=(args.beta1, args.beta2),
                weight_decay=args.wd,
                eps=args.eps,
            )
optimizer.add_p…