code: def quantize_state_dict(
      module: nn.Module,
      table_name_to_quantized_weights: Dict[str, Tuple[Tensor, Tensor]],
      table_name_to_data_type: Dict[str, DataType],
      table_name_to_pruning_indices_mapping: Optional[
          Dict[str, Tensor]
      ] = None,
  ) -> torch.device:
      device = torch.device("cpu")
      if not table_name_to_pruning_indices_mapping:
          table_name_to_pruning_indices_mapping = {}
      for key, tensor in module.state_dict().items():
          # Extract table name from state dict key.
          # e.g. ebc.embedding_bags.t1.weight
          splits = key.split(".")
          assert splits[-1] == "weight"
          table_name = splits[-2]
          data_type = table_name_to_data_type[table_name]
          num_rows = tensor.shape[0]
          pruning_indices_mapping: Optional[Tensor] = None
          if table_name in table_name_to_pruning_indices_mapping:
              pruning_indices_mapping = (
                  table_name_to_pruning_indices_mapping[table_name]
              )
          if pruning_indices_mapping is not None:
              num_rows = pruned_num_embeddings(pruning_indices_mapping)
  
          device = tensor.device
          num_bits = DATA_TYPE_NUM_BITS[data_type]
  
          if tensor.is_meta:
              quant_weight = torch.empty(
                  (num_rows, (tensor.shape[1] * num_bits) // 8),
                  device="meta",
                  dtype=torch.uint8,
              )
              if (
                  data_type == DataType.INT8
                  or data_type == DataType.INT4
                  or data_type == DataType.INT2
              ):
                  scale_shift = torch.empty(
                      (num_rows, 4),
                      device="meta",
                      dtype=torch.uint8,
                  )
              else:
                  scale_shift = None
          else:
              if pruning_indices_mapping is not None:
                  rows_mask = pruning_indices_mapping.gt(-1)
                  tensor = tensor[rows_mask, :]
  
              if (
                  tensor.dtype == torch.float
                  or tensor.dtype == torch.float16
              ):
                  if data_type == DataType.FP16:
                      if tensor.dtype == torch.float:
                          tensor = tensor.half()
                      quant_res = tensor.view(torch.uint8)
                  else:
                      quant_res = torch.ops.fbgemm.FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
                          tensor, num_bits
                      )
              else:
                  raise Exception("Unsupported dtype: {tensor.dtype}")
              if (
                  data_type == DataType.INT8
                  or data_type == DataType.INT4
                  or data_type == DataType.INT2
              ):
                  quant_weight, scale_shift = (
                      quant_res[:, :-4],
                      quant_res[:, -4:],
                  )
              else:
                  quant_weight, scale_shift = quant_res, None
          table_name_to_quantized_weights[table_name] = (
              quant_weight,
              scale_shift,
          )
      return device
  
imports: import copy
  import itertools
  from collections import defaultdict
  from typing import Callable, cast, Dict, List, Optional, Tuple, Type, Union
  import torch
  import torch.nn
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import EmbeddingLocation, IntNBitTableBatchedEmbeddingBagsCodegen, PoolingMode
  from torch import Tensor
  from torchrec.distributed.utils import none_throws
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, DATA_TYPE_NUM_BITS, data_type_to_sparse_type, DataType, dtype_to_data_type, EmbeddingBagConfig, EmbeddingConfig, pooling_type_to_pooling_mode, PoolingType, QuantConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingBagCollectionInterface, EmbeddingCollection, EmbeddingCollectionInterface, get_embedding_names_by_table
  from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
  from torchrec.tensor_types import UInt2Tensor, UInt4Tensor
  from torchrec.types import ModuleNoCopyMixin
label_desc: Use image node_id 8 for calling a global function with example usage: quantize_state_dict(module, table_name_to_quantized_weights, table_name_to_data_type, table_name_to_pruning_indices_mapping) and returns: device