warm up LR causes crash
My resnet encoder and transformer decoder are not training well.  So trying all kinds of stuff.
Latest attempt to improve is to use a warmup learning rate as described here:
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/docs/source/optimizers.rst
My code is an exact copy:
    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, opt_closure):
        if self.trainer.global_step < 500:
            lr_scale = min(1., float(self.trainer.global_step + 1) / 500.)
            for pg in optimizer.param_groups:
                pg['lr'] = lr_scale * self.hparams.learning_rate

                print(f"lr={pg['lr']}")
        optimizer.step()
        optimizer.zero_grad()

The crash is as follows:
Epoch 1:   0%|▍                                                                                                                     | 15/3544 [00:27<1:48:33,  1.85s/it, loss=nan, v_num=19]Traceback (most recent call last):
  File "kiss_transformer.py", line 539, in <module>
    trainer.fit(model, train_loader)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/states.py", line 34, in wrapped_fn
    result = fn(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1023, in fit
    self.accelerator_backend.train(model, nprocs=self.num_processes)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 42, in train
    mp.spawn(self.ddp_train, nprocs=nprocs, args=(self.mp_queue, model,))
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 158, in start_processes
    while not context.join():
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 119, in join
    raise Exception(msg)
Exception: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/ddp_spawn_backend.py", line 154, in ddp_train
    results = self.trainer.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1211, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 393, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 490, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 887, in run_training_batch
    grad_norm_dic = self.run_batch_backward_pass(split_batch, batch_idx, opt_idx, optimizer)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 948, in run_batch_backward_pass
    self.call_optimizer_step(optimizer, opt_idx, batch_idx, split_batch)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 986, in call_optimizer_step
    using_native_amp=native_amp)
TypeError: optimizer_step() got an unexpected keyword argument 'using_native_amp'


I have tried both these versions and same crash:
#RUN pip install pytorch-lightning==0.8.5
RUN pip install pytorch-lightning==0.9.0rc12

Not that I believe these things are related, but my code has the following:
        trainer = pl.Trainer(gpus=[0, 1],  accumulate_grad_batches=16, callbacks=[lr_logger]) #  
        trainer.fit(model, train_loader)