Replace torch.cuda.amp.autocast with torch.amp.autocast('cuda',...)

Replace torch.cuda.amp.autocast with torch.amp.autocast('cuda',...) #71