How to correctly initialize latent vector parameters that have size dependent on training data size using CLI?


May I ask how do you correctly create a set of latents for each sample in the training dataset? I.e., suppose you would like to have optimizable latent codes for each of the frame. The total number of those latents is then decided by the number of data samples.

I was unable to do it in __init__of LightningModule, as at that point DataModule is not loaded yet. I have manged to put it at the beginning of configure_optimizers, however it causes problem when trying to load a trained ckpt for prediction or resuming training, as the loading of weights seem to happen after __init__ and DataModule loading, but before configure_optimizers.

May I ask if there is a way to do it, or it can’t be achieved using lighting CLI?