[Optim] Add unittests for PyTorch native AMP support

PyTorch native AMP (i.e., `torch.cuda.amp`) interacts with the optimizers and expects certain behaviors out of the optimizers. We should add unit tests to make sure that those interactions are fine.