Corrdiff error

Hi,
I am trying to run corrdiff using the checkpoint available here:
https://catalog.ngc.nvidia.com/orgs/nvidia/teams/modulus/models/corrdiff_inference_package

I have installed the package as instructed with

git checkout ee6b3ac0d4c32bca1c81f401d2bfd461887eb9a3

Then when I run the generate script it seems to work to nearly the end of the inference and I get the following error
Traceback (most recent call last):
File “/home/whittaker/modulus/examples/generative/corrdiff/generate.py”, line 302, in main
generate_and_save(
File “/home/whittaker/modulus/examples/generative/corrdiff/generate.py”, line 388, in generate_and_save
image_out = generate_fn(image_lr)
^^^^^^^^^^^^^^^^^^^^^
File “/home/whittaker/modulus/examples/generative/corrdiff/generate.py”, line 239, in generate_fn
image_res = generate(
^^^^^^^^^
File “/home/whittaker/modulus/examples/generative/corrdiff/generate.py”, line 479, in generate
latents = rnd.randn(
^^^^^^^^^^
File “/Data/conda/envs/modu/lib/python3.11/site-packages/modulus/utils/generative/utils.py”, line 76, in randn
raise ValueError(
ValueError: Expected first dimension of size 4, got 5

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Has anyone managed to get it to run?