code: def _unwrap_kjt(
      features: KeyedJaggedTensor,
  ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
      if features.device().type == "cuda":
          return (
              features.values().int(),
              features.offsets().int(),
              features.weights_or_none(),
          )
      else:
          return (
              features.values(),
              features.offsets(),
              features.weights_or_none(),
          )
  
imports: import copy
  import logging
  from typing import Any, Dict, Iterator, List, Optional, Tuple
  import torch
  import torch.distributed
  from fbgemm_gpu.split_embedding_configs import SparseType
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import EmbeddingLocation, IntNBitTableBatchedEmbeddingBagsCodegen, PoolingMode, rounded_row_size_in_bytes
  from torchrec.distributed.batched_embedding_kernel import BaseBatchedEmbedding, BaseBatchedEmbeddingBag, BatchedDenseEmbedding, BatchedDenseEmbeddingBag
  from torchrec.distributed.embedding_kernel import BaseEmbedding
  from torchrec.distributed.embedding_types import compute_kernel_to_embedding_location, GroupedEmbeddingConfig
  from torchrec.distributed.fused_params import is_fused_param_quant_state_dict_split_scale_bias, is_fused_param_register_tbe, tbe_fused_params, TBEToRegisterMixIn
  from torchrec.distributed.utils import append_prefix
  from torchrec.modules.embedding_configs import DATA_TYPE_NUM_BITS, data_type_to_sparse_type, DataType, dtype_to_data_type
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 3 for calling a global function with example usage: _unwrap_kjt(features) and returns: features,  features,  features, features,  features,  features