Fix the warning: ``` FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. ``` by replacing everywhere in the codebase: ```diff - torch.cuda.amp.autocast(... + torch.amp.autocast("cuda", ... ```