Add another parameter in `Trainer.train()` to enable `torch.amp`.
Add another parameter in
Trainer.train()to enabletorch.amp.