MSSSIM VAE is not working with 3D inputs

[ravih18]

Hello @clementchadebec

MSSSIM VAE model returns an error when using 3D images for training.

Indeed the MSSSIM implementation in benchmark_VAE/src/pythae/models/msssim_vae/msssim_vae_utils.py only works for 2 images.

I found the following implementation that seems to work with 3D images: https://github.com/VainF/pytorch-msssim.

I can make a PR to add it if you think it is a good idea.

Otherwise I can see if I can generalize the current implementation for 3D images !

Let me know what do you think of it.

Ravi

[clementchadebec]

Hello @ravih18,

Thank you for opening this issue. You are absolutely right, the current implementation of VAE-MSSSIM only supports 2D images. Having an implementation that supports 3D images would indeed be a very nice feature.

As for the PR, sure! Thanks for suggesting it. Feel free to send a PR as this feature may be interesting for other users as well! I would be happy to look into it :) Feel free to contact me if you have any questions about integrating this feature into the code base.

Best,

Clément

[clementchadebec]

Completed in #69