[benchmarks] Fix AMP setup for torchbench models.

[ysiraichi]

Fix: #6556 (and, possibly #6833)

This PR fixes the benchmarks script when running with AMP. Previously, we were calling torch.amp.autocast(..., device_type="xla") for both XLA:CUDA and XLA:TPU. However, we should be using torch.cuda.amp.autocast for XLA:CUDA (see this for more details).

Context: after #6518, Super_Slomo inference started being run using AMP. However, due to #6511, that PR tried to mimic torch_xla.amp.autocast behavior, using torch.amp.autocast.

cc @miladm @JackCaoG @vanbasten23 @zpcore

[ysiraichi]

Confirmed it also fixes #6833.

[JackCaoG]

Hmm according to https://github.com/pytorch/xla/blob/master/docs/amp.md we should be able to use autocast for both TPU and GPU, is that no longer the case?

[ysiraichi]

That document is correct. Problem is that I didn't notice XLA:CUDA is supposed to run with CUDA autocast, i.e. torch.amp.autocast("cuda"). Instead, I was running both XLA:CUDA and XLA:TPU with XLA autocast, i.e. torch.amp.autocast("xla"). This behavior is implemented in torch_xla.amp.autocast. However, since it currently doesn't work (#6511) with dynamo, I was using torch.amp.autocast directly.

[vanbasten23]

I guess torch.cuda.amp.autocast is the same as torch.amp.autocast("cuda")?

[vanbasten23]

do you need to set kwargs["device_type"] = "xla" for XLA:GPU case?

[ysiraichi]

Not really. torch.cuda.amp.autocast already does that.