Model export

[ctr26]

torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/home/ctr26/gdrive/+projects/2023_shape_embed/env/lib/python3.9/site-packages/pythae/models/vq_vae/vq_vae_model.py", line 74
    def forward(self, inputs: BaseDataset, **kwargs):
                                            ~~~~~~~ <--- HERE
        """
        The VAE model

I'd be useful to have the code base setup such that all the models have fixed inputs so that they can be compiled and exported

Thanks

Craig