def METHOD_NAME(self) -> torch.autocast:
    # the dtype could be automatically inferred but we need to manually set it due to a bug upstream
    # https://github.com/pytorch/pytorch/issues/67233
    return torch.autocast(self.device, dtype=self._desired_input_dtype)