mmacosha's picture
add mnist_vae/ (checkpoint + code)
5e3f03f verified
Raw
History Blame Contribute Delete
396 Bytes
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"],
)