lattice_point and walkthrough doesn't work except z_dim=2

[shun60s]

It seems that lattice_point and walkthrough just works when z_dim=2.
For example, z_dim=10, following error rises.

python3\lib\site-packages\torch\nn\functional.py", line 1692, in linear
output = input.matmul(weight.t())
RuntimeError: mat1 and mat2 shapes cannot be multiplied (625x2 and 10x200)