I want loaded AutoencoderKL objects to provide a pure latent-to-image decode API: decode(z…

I want loaded AutoencoderKL objects to provide a pure latent-to-image decode API: decode(z…: a task in MiMo-V2.6-RL-harbor-code: MiMo-V2.6-RL Code (Harbor) (Harbor dataset). When I create AutoencoderKL(latent channels=4, out channels=3, block out channels=(32,), layers per block=1, norm num…

The task

When I create `AutoencoderKL(latent_channels=4, out_channels=3, block_out_channels=(32,), layers_per_block=1, norm_num_groups=32, use_post_quant_conv=True)`, zero all decoder and post-quantization parameters, and call `decode(torch.zeros(2, 4, 8, 8))`, I should get an `AutoencoderKLOutput` whose `.sample` is a zero…

Part of FineEnvs/MiMo-V2.6-RL-harbor-code.