Keyword argument in `torch.autocast()`

[smly]

This PR resolves an error caused by a change in the keyword arguments in torch.autocast().

The order of keyword arguments in torch.autocast() has changed, so you will get an error if you use the latest version of PyTorch. torch.autocast()’s second keyword argument is enabled= until PyTorch 1.10. Since 1.11, dtype= is the second keyword argument.

Reference:

• https://pytorch.org/docs/1.10/amp.html?highlight=autocast#torch.autocast
• https://pytorch.org/docs/stable/amp.html?highlight=autocast#torch.autocast