[BUG] Update usage of deprecated amp APIs

[svekars]

Add Link

https://pytorch.org/docs/stable/amp.html

Describe the bug

See the link above for the list of deprecated APIs.

We need to fix:

FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.

And same for:

torch.cuda.amp.autocast(args...) and torch.cpu.amp.autocast(args...) will be deprecated. Please use torch.autocast("cuda", args...) or torch.autocast("cpu", args...) instead.

torch.cuda.amp.GradScaler(args...) and torch.cpu.amp.GradScaler(args...) will be deprecated. Please use torch.GradScaler("cuda", args...) or torch.GradScaler("cpu", args...) instead.

Describe your environment

• PyTorch 2.6 or later

[nirajkamal]

Hi @svekars, I am working on this issue, can you please assign it to me?

[nirajkamal]

Thank you @svekars, I think the branch is merged with changes for this issue.

Also do you think it's more appropriate to use torch.amp.autocast(...) instead of torch.autocast(....) / torch.amp.GradScaler(...) instead of torch.GradScaler(...)?
in the docs the usage is not consistent.