src/transformers/trainer.py/Trainer/autocast_smart_context_manager
class Trainer:    def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True):
        """
        A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired
        arguments, depending on the situation.
        """
        if self.use_cuda_amp or self.use_cpu_amp:
            if is_torch_greater_or_equal_than_1_10:
                ctx_manager = (
                    torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
                    if self.use_cpu_amp
                    else torch.cuda.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
                )
            else:
                ctx_manager = torch.cuda.amp.autocast()
        else:
            ctx_manager = contextlib.nullcontext() if sys.version_info >= (3, 7) else contextlib.suppress()

        return ctx_manager