Can you give a bit more detail on why this is needed?

Otherwise, we are calling nullcontext(dtype=...), here:

Which ends up raising an error.

Otherwise, we are calling nullcontext(dtype=...), here:

xla/benchmarks/benchmark_model.py

Lines 155 to 159 in 3e5ca3c

	def train(self, inputs, collect_full_output=False):
	self._optimizer_zero_grad()
	with self.autocast(**self.autocast_kwargs):
	pred = self.module(*inputs)
	loss = self.compute_loss(pred)

Which ends up raising an error.

Not sure if I follow. With the change, for amp case, we'll call nullcontext(dtype=...) with non-empy kwargs.., wouldn't nullcontext complain too?

As far as I understand it, for AMP case, self.autocast will be either torch.amp.autocast or torch.cuda.autocast. Am I missing something?