code: def _reserve_kjt_storage(
      topology: Topology,
      batch_size: int,
      batch_inputs: List[float],
      input_data_type_size: int,
      multiplier: int,
  ) -> Storage:
      kjt_size = (
          math.ceil(sum(batch_inputs) * float(input_data_type_size))
          * multiplier
      )
  
      kjt_storage = Storage(
          hbm=kjt_size if topology.compute_device == "cuda" else 0,
          ddr=kjt_size
          if topology.compute_device in {"cpu", "mtia"}
          else 0,
      )
  
      for device in topology.devices:
          device.storage -= kjt_storage
  
      return kjt_storage
  
imports: import copy
  import logging
  import math
  from typing import Dict, List, Optional, Set, Tuple
  from torch import nn
  from torchrec.distributed.planner.constants import BIGINT_DTYPE, POOLING_FACTOR
  from torchrec.distributed.planner.types import ParameterConstraints, PlannerError, PlannerErrorType, Storage, StorageReservation, Topology
  from torchrec.distributed.planner.utils import sharder_name, storage_repr_in_gb
  from torchrec.distributed.types import get_tensor_size_bytes, ModuleSharder
label_desc: Use image node_id 4 for calling a global function with example usage: _reserve_kjt_storage(topology, batch_size, batch_inputs, input_data_type_size, multiplier) and returns: kjt_storage