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?