import pickle

from vae import VAE

with open("mnist_vae.pkl", "rb") as f:
    bundle = pickle.load(f)
meta = bundle["meta"]
params = bundle["params"]

model = VAE(
    latent_dim=meta["latent_dim"],
    in_channels=meta["in_channels"],
    base_channels=meta["base_channels"],
    channel_mults=tuple(meta["channel_mults"]),
    image_size=meta["image_size"],
    obs_sigma=meta["obs_sigma"],
)