import numpy as np
import torch
def produce_embeddings(images_factor_dataset, encoder_pytorch,z_dim,  device = "cpu"):
    latent = np.zeros(images_factor_dataset.factors_shape + (z_dim,))
    for num_factor1, factor1 in enumerate(images_factor_dataset.factor_values_list[0]):
        for num_factor2, factor2 in enumerate(images_factor_dataset.factor_values_list[1]):
            print("Encoding image ({},{})".format(factor1, factor2))
            img = images_factor_dataset.images[num_factor1, num_factor2, ...]
            torch.from_numpy(img).permute(-1, 0, 1).float()
            z = encoder_pytorch(torch.from_numpy(img).permute(-1, 0, 1).float().to(device))  # .numpy()
            latent[num_factor1, num_factor2] = z.detach().to("cpu").numpy()
    return latent