when unloading vae during TI training, generating images after
training will error out