from transformers import AutoModel

MODEL_ID = "pszmk/mnist-vae-latent2"
REVISION = "main"

model = AutoModel.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    trust_remote_code=True,
)

print(model.__class__.__name__)