Calling a different method post training

[shenoynikhil]

I have a variational autoencoder from which I want to generate images. Post training, my VAE model class has a separate function for generating images from gaussian noise,

from flax import linen as nn

class VAE(nn.Module):
   def generate_samples(self, num_samples):

How do I call this method with say a state that I have post training?

[IvyZX]

You can directly access the parameters from state.params. Then you can call your own method with it.

For example for nn.Module's apply(), you can run inference like: model.apply({'params': state.params}, x)