code: def _extract_device_ids(
      forward_fn: Callable,
      saved_layer: Dict[device, Tuple[Tensor, ...]],
      device_ids: Union[None, List[int]],
  ) -> Union[None, List[int]]:
      r"""
      A helper function to extract device_ids from `forward_function` in case it is
      provided as part of a `DataParallel` model or if is accessible from
      `forward_fn`.
      In case input device_ids is not None, this function returns that value.
      """
      # Multiple devices / keys implies a DataParallel model, so we look for
      # device IDs if given or available from forward function
      # (DataParallel model object).
      if len(saved_layer) > 1 and device_ids is None:
          if (
              hasattr(forward_fn, "device_ids")
              and cast(Any, forward_fn).device_ids is not None
          ):
              device_ids = cast(Any, forward_fn).device_ids
          else:
              raise AssertionError(
                  "Layer tensors are saved on multiple devices, however unable to access"
                  " device ID list from the `forward_fn`. Device ID list must be"
                  " accessible from `forward_fn`. For example, they can be retrieved"
                  " if `forward_fn` is a model of type `DataParallel`. It is used"
                  " for identifying device batch ordering."
              )
      return device_ids
  
imports: import threading
  import typing
  import warnings
  from typing import Any, Callable, Dict, List, Tuple, Union, cast
  import torch
  from torch import Tensor, device
  from torch.nn import Module
  from captum.attr._utils.batching import _reduce_list, _sort_key_list
  from captum.attr._utils.common import _run_forward, _verify_select_column
  from captum.attr._utils.typing import Literal, TargetType, TensorOrTupleOfTensorsGeneric
label_desc: Use image node_id 10 for calling a global function with example usage: _extract_device_ids(forward_fn, saved_layer, device_ids) and returns: device_ids