def METHOD_NAME(self) -> None:
    B = 2
    D = 8
    num_dense_features = 100
    eb1_config = EmbeddingBagConfig(
        name="t1", embedding_dim=D, num_embeddings=100, feature_names=["f1", "f3"]
    )
    eb2_config = EmbeddingBagConfig(
        name="t2",
        embedding_dim=D,
        num_embeddings=100,
        feature_names=["f2"],
    )
    ebc = EmbeddingBagCollection(tables=[eb1_config, eb2_config])
    deepfm_nn = SimpleDeepFMNN(
        num_dense_features=num_dense_features,
        embedding_bag_collection=ebc,
        hidden_layer_size=20,
        deep_fm_dimension=5,
    )
    features = torch.rand((B, num_dense_features))
    sparse_features = KeyedJaggedTensor.from_offsets_sync(
        keys=["f1", "f3", "f2"],
        values=torch.tensor([1, 2, 4, 5, 4, 3, 2, 9, 1, 2, 3]),
        offsets=torch.tensor([0, 2, 4, 6, 8, 10, 11]),
    )
    deepfm_nn(
        dense_features=features,
        sparse_features=sparse_features,
    )
    gm = symbolic_trace(deepfm_nn)
    scripted_gm = torch.jit.script(gm)
    logits = scripted_gm(features, sparse_features)
    self.assertEqual(logits.size(), (B, 1))