🐛[BUG]: CorrDiff - Inference using Pretrained Model Can't Run Succussfully, which Lacks Necessary Arugments

[windsoryin]

Version

0.6.0

On which installation method(s) does this occur?

Docker, Pip

Describe the issue

The inference of CorrDiff, using Nvidia Inference Package, i.e., a pretrained model and sample data, will raise errors, same in issue: #542 .
This is due to the wrong definition in config_generate.yaml.
in_channels: [0, 1, 2, 3, 4, 9, 10, 11, 12, 17, 18, 19]
out_channels: [0, 17, 18, 19]

According to the config file, the input data has 12 channels, but the pretrained model needs 16 input channels.
After i changed the config: in_channels=[ 1, 2, 3, 4,5,6,7,8, 9, 10, 11, 12,13,14,15,16], which includes four pressure levels (500 hPa, 700 hPa, 850 hPa, and 925 hPa) with four variables: temperature, 2 wind vectors, and geopotential Height.
But the prediction results are not satisfactory at all.

So what's the exact in_channels option used in the pre-trainning process? And the pre-trained model is not 20 in_channels as is stated in the paper。The instructions on inference using pre-trained could be better if necessary arguments and description are added.

Minimum reproducible example

following the instructions of CorrDiff Inference Package website.
Getting started
Follow these steps to set up all needed to run the CorrDiff inference:

Clone the NVIDIA Modulus package, check out a specific hash, and install it:
git clone https://github.com/NVIDIA/modulus
cd modulus
git checkout ee6b3ac0d4c32bca1c81f401d2bfd461887eb9a3
pip install .
Get this package using direct download or via wget:
wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/modulus/corrdiff_inference_package/versions/1/zip -O corrdiff_inference_package_1.zip
Unzip the contents in the same directory:
unzip corrdiff_inference_package_1.zip
unzip corrdiff_inference_package.zip
Assuming you have both the Modulus and the unzipped CorrDiff inference packages in the <root_dir> directory, go to the directory that includes the CorrDiff inference script:
cd <root_dir>/modulus/examples/generative/corrdiff
Run the inference:
python generate.py dataset.data_path=<root_dir>/corrdiff_inference_package/dataset/2023-01-24-cwb-4years_5times.zarr res_ckpt_filename=<root_dir>/corrdiff_inference_package/checkpoints/diffusion.mdlus reg_ckpt_filename=<root_dir>/corrdiff_inference_package/checkpoints/regression.mdlus seed_batch_size=5 use_torch_compile=false

Relevant log output

No response

Environment details

No response