problem with forward_noise = self._distribution.sample(x.shape[0]).view(x.shape)

[laurentbenaroya]

Thank you for this great work. Though, I couldn't make it work using poetry, it gave me problems with cuda and get-poetry.py is outdated. So I used a conda environnement and things went fine. I just tried run the command python diffflow/main.py trainer.epochs=1 data.dataset=tree and I get the following error : File "/home/xxx/DiffFlow/diffflow/networks/base_model.py", line 26, in forward_step forward_noise = self._distribution.sample(x.shape[0]).view(x.shape) RuntimeError: shape '[2]' is invalid for input of size 4 I also tried with data.dataset=MNIST, but I got a similar error. In the base_model.py, in the constructor : self._distribution = StandardNormal([np.prod(data_shape)]) Could anyone help me?