def test_image_embeddings_sanity(image_params_dict):
    class OldImageFeatureEmbeddings(TransformerModule, FromParams):
        """Construct the embeddings from image, spatial location (omit now) and
        token_type embeddings.
        """

        def __init__(self, feature_size: int, embedding_size: int, dropout: float = 0.0):
            super().__init__()

            self.image_embeddings = torch.nn.Linear(feature_size, embedding_size)
            self.image_location_embeddings = torch.nn.Linear(4, embedding_size, bias=False)
            self.layer_norm = torch.nn.LayerNorm(embedding_size, eps=1e-12)
            self.dropout = torch.nn.Dropout(dropout)

        def forward(self, image_feature: torch.Tensor, image_location: torch.Tensor):
            img_embeddings = self.image_embeddings(image_feature)
            loc_embeddings = self.image_location_embeddings(image_location)
            embeddings = self.layer_norm(img_embeddings + loc_embeddings)
            embeddings = self.dropout(embeddings)

            return embeddings

    torch.manual_seed(23)
    old = OldImageFeatureEmbeddings(**image_params_dict)
    torch.manual_seed(23)
    now = ImageFeatureEmbeddings(**image_params_dict)

    batch_size = 2

    image_feature = torch.randn(batch_size, image_params_dict["feature_size"])
    image_location = torch.randn(batch_size, 4)

    torch.manual_seed(23)
    old_output = old(image_feature, image_location)
    torch.manual_seed(23)
    now_output = now(image_feature, image_location)

    assert_allclose(old_output, now_output)
Construct the embeddings from image, spatial location (omit now) and
token_type embeddings.
tests/modules/transformer/transformer_embeddings_test.py