Function optimize_fn

Optimizes with warmup and gradient clipping (disabled if negative).

(optimizer, params, step, lr=config.optim.lr,
                  warmup=config.optim.warmup,
                  grad_clip=config.optim.grad_clip)