- Move the estimation logic to utility functions - Estimate memory _within_ the encode and decode methods, ensuring we _always_ estimate working memory when running a VAE