code: def fuse_embedding_optimizer(
      model: nn.Module,
      optimizer_type: Type[torch.optim.Optimizer],
      optimizer_kwargs: Dict[str, Any],
      device: torch.device,
      location: Optional[EmbeddingLocation] = None,
  ) -> nn.Module:
      """
      Recursively replaces EmbeddingBagCollection and EmbeddingCollection with
      FusedEmbeddingBagCollection and FusedEmbeddingCollection in a model subtree.
      The fused modules will be initialized using the passed in optimizer parameters, and model location.
  
      Args:
          model: (nn.Module):
          optimizer_type: (Type[torch.optim.Optimizer]):
          optimizer_kwargs: (Dict[Str, Any]):
          device (Optional[torch.device]):
          location: (Optional[EmbeddingLocation]): GPU location placement
      Returns
          nn.Module: input nn.Module with Fused Embedding Modules
  
      Example::
          ebc = EmbeddingBagCollection()
          my_model = ExampleModel(ebc)
          my_model = fused_embedding_optimizer(my_model, optimizer_type=torch.optim.SGD, optimizer_kwargs={"lr": .01})
          kjt = KeyedJaggedTensor()
          output = my_model(kjt)
      """
      # Replace all EBCs and ECs in a with a corresponding FusedEmbeddingModule.
  
      # check if top-level module is EBC/EC
      if isinstance(model, EmbeddingBagCollection):
          return FusedEmbeddingBagCollection(
              model.embedding_bag_configs(),
              optimizer_type=optimizer_type,
              optimizer_kwargs=optimizer_kwargs,
              device=device,
              location=location,
          )
      if isinstance(model, EmbeddingCollection):
          return FusedEmbeddingCollection(
              model.embedding_configs(),
              optimizer_type=optimizer_type,
              optimizer_kwargs=optimizer_kwargs,
              device=device,
              location=location,
          )
  
      def replace(_model: nn.Module) -> None:
          for child_name, child in _model.named_children():
              if isinstance(child, EmbeddingBagCollection):
                  setattr(
                      _model,
                      child_name,
                      FusedEmbeddingBagCollection(
                          tables=child.embedding_bag_configs(),
                          optimizer_type=optimizer_type,
                          optimizer_kwargs=optimizer_kwargs,
                          device=device,
                          location=location,
                      ),
                  )
              elif isinstance(child, EmbeddingCollection):
                  setattr(
                      _model,
                      child_name,
                      FusedEmbeddingCollection(
                          tables=child.embedding_configs(),
                          optimizer_type=optimizer_type,
                          optimizer_kwargs=optimizer_kwargs,
                          device=device,
                          location=location,
                      ),
                  )
              else:
                  replace(child)
  
      replace(model)
      return model
  
imports: import copy
  import itertools
  from collections import defaultdict
  from typing import Any, cast, Dict, Iterator, List, Optional, Set, Tuple, Type
  import torch
  import torch.nn
  import torchrec.optim
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from fbgemm_gpu.split_table_batched_embeddings_ops_training import ComputeDevice, EmbeddingLocation, SplitTableBatchedEmbeddingBagsCodegen
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, DataType, EmbeddingBagConfig, EmbeddingConfig, pooling_type_to_pooling_mode, PoolingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingBagCollectionInterface, EmbeddingCollection, EmbeddingCollectionInterface, get_embedding_names_by_table
  from torchrec.optim.fused import FusedOptimizer, FusedOptimizerModule
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 2 for calling a global function with example usage: fuse_embedding_optimizer(model, optimizer_type, optimizer_kwargs, device, location) and returns: model, FusedEmbeddingBagCollection, FusedEmbeddingCollection