[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