code: def _get_adjusted_batch_keys(
      batch_state_dict: Dict[str, Tensor],
      input_transform: Optional[InputTransform],
      outcome_transform: Optional[OutcomeTransform] = None,
  ) -> Tuple[Set[str], Set[str]]:
      r"""Group the keys based on whether the value requires batch shape changes.
  
      Args:
          batch_state_dict: The state dict of the batch model.
          input_transform: The input transform.
          outcome_transform: The outcome transform.
  
      Returns:
          A two-element tuple containing:
              - The keys of the parameters/buffers that require a batch shape adjustment.
              - The keys of the parameters/buffers that do not require a batch shape
                  adjustment.
      """
      # These are the names of the params/buffers that need their batch shape adjusted.
      adjusted_batch_keys = {
          n for n, p in batch_state_dict.items() if len(p.shape) > 0
      }
      # Don't modify transform buffers, so add them to non-adjusted set and remove
      # them from tensors.
      for transform, transform_type in [
          (input_transform, "input_transform."),
          (outcome_transform, "outcome_transform."),
      ]:
          if transform is not None:
              transform_keys = {
                  transform_type + n
                  for n, p in transform.state_dict().items()
              }
              adjusted_batch_keys = adjusted_batch_keys - transform_keys
      # These are the names of the parameters/buffers that don't need their
      # batch shape adjusted.
      non_adjusted_batch_keys = (
          set(batch_state_dict) - adjusted_batch_keys
      )
      return adjusted_batch_keys, non_adjusted_batch_keys
  
imports: from __future__ import annotations
  from copy import deepcopy
  from typing import Dict, Optional, Set, Tuple
  import torch
  from botorch.exceptions import UnsupportedError
  from botorch.models.gp_regression import HeteroskedasticSingleTaskGP
  from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
  from botorch.models.gp_regression_mixed import MixedSingleTaskGP
  from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
  from torch import Tensor
  from torch.nn import Module
label_desc: Use image node_id 9 for calling a global function with example usage: _get_adjusted_batch_keys(batch_state_dict, input_transform, outcome_transform) and returns: adjusted_batch_keys,  non_adjusted_batch_keys