code: def test_forward_with_state_dict(
      self,
      device: torch.device,
  ) -> None:
      # this tests the common components of an EmbeddingCollection, namely
      # lookup with multiple features, shared features,
      # calling state_dict and calling load_state_dict
      e1_config = EmbeddingConfig(
          name="t1",
          embedding_dim=4,
          num_embeddings=2,
          feature_names=["f1", "f1_1", "f_shared"],
      )
      e2_config = EmbeddingConfig(
          name="t2",
          embedding_dim=4,
          num_embeddings=2,
          feature_names=["f2", "f_shared"],
      )
      e3_config = EmbeddingConfig(
          name="t3",
          embedding_dim=4,
          num_embeddings=2,
          feature_names=["f3"],
      )
  
      ec = FusedEmbeddingCollection(
          tables=[e1_config, e2_config, e3_config],
          optimizer_type=torch.optim.SGD,
          optimizer_kwargs={"lr": 0.02},
          device=device,
      )
  
      non_fused_ec = EmbeddingCollection(
          tables=[e1_config, e2_config, e3_config],
          device=device,
      )
  
      ec.load_state_dict(
          OrderedDict(
              [
                  (
                      "embeddings.t1.weight",
                      torch.Tensor([[1, 1, 1, 1], [2, 2, 2, 2]]).to(
                          device
                      ),
                  ),
                  (
                      "embeddings.t2.weight",
                      torch.Tensor([[4, 4, 4, 4], [8, 8, 8, 8]]).to(
                          device
                      ),
                  ),
                  (
                      "embeddings.t3.weight",
                      torch.Tensor(
                          [[16, 16, 16, 16], [32, 32, 32, 32]]
                      ).to(device),
                  ),
              ]
          )
      )
  
      non_fused_ec.load_state_dict(ec.state_dict())
  
      #    0       1        2  <-- batch
      # f1   [0,1] []    [0]
      # f1_1   [0] [1]    [1,0]
      # f2   [0]    [1]    [0,1]
      # f3   []    []    [0]
      # f_shared   [0]    [1]    [0,1]
      # ^
      # feature
      features = KeyedJaggedTensor.from_lengths_sync(
          keys=["f1", "f1_1", "f2", "f3", "f_shared"],
          values=torch.tensor(
              [0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1]
          ),
          lengths=torch.tensor(
              [2, 0, 1, 1, 1, 2, 1, 1, 2, 0, 0, 1, 1, 1, 2]
          ),
      ).to(device)
  
      sequential_embeddings = ec(features)
      self.assertEquals(
          set(sequential_embeddings.keys()),
          {"f1", "f1_1", "f_shared@t1", "f2", "f_shared@t2", "f3"},
      )
  
      self._assert_dense_list_equality(
          sequential_embeddings["f1"].to_dense(),
          [
              torch.Tensor(
                  [[1.0, 1.0, 1.0, 1.0], [2.0, 2.0, 2.0, 2.0]]
              ).to(device),
              torch.empty(0, 4).to(device),
              torch.Tensor([[1.0, 1.0, 1.0, 1.0]]).to(device),
          ],
      )
  
      self._assert_dense_list_equality(
          sequential_embeddings["f1_1"].to_dense(),
          [
              torch.Tensor([[1.0, 1.0, 1.0, 1.0]]).to(device),
              torch.Tensor([[2.0, 2.0, 2.0, 2.0]]).to(device),
              torch.Tensor(
                  [[2.0, 2.0, 2.0, 2.0], [1.0, 1.0, 1.0, 1.0]]
              ).to(device),
          ],
      )
  
      non_fused_sequential_embeddings = non_fused_ec(features)
  
      for key in sequential_embeddings:
          self._assert_dense_list_equality(
              sequential_embeddings[key].to_dense(),
              non_fused_sequential_embeddings[key].to_dense(),
          )
  
imports: import unittest
  from collections import OrderedDict
  from typing import Any, Dict, List, Optional, Tuple, Type
  import hypothesis.strategies
  import torch
  import torch.fx
  import torchrec
  from fbgemm_gpu.split_table_batched_embeddings_ops_training import EmbeddingLocation
  from hypothesis import given, settings
  from torchrec.modules.embedding_configs import EmbeddingBagConfig, EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingCollection
  from torchrec.modules.fused_embedding_modules import fuse_embedding_optimizer, FusedEmbeddingBagCollection, FusedEmbeddingCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the FusedEmbeddingCollectionTest obj's underlying member method code with example usage: obj.test_forward_with_state_dict(device) without return types