Fix missing num_steps parameter for stochastic sampler

[younes-abid]

• Add missing num_steps=cfg.sampler.num_steps parameter to stochastic_sampler partial() call
• This bug caused stochastic sampler to always use default 18 steps instead of configured num_steps
• Fixes inconsistency with deterministic sampler which correctly passes num_steps parameter
• Improves performance when using different diffusion num_steps

PhysicsNeMo Pull Request

Description

This PR fixes a bug in the CorrDiff generation pipeline where the stochastic sampler was ignoring the configured num_steps
parameter and always using the default value of 18 steps, causing performance issues.

Problem: In generate.py, the partial() call for stochastic_sampler was missing the num_steps parameter, while the deterministic sampler correctly included it.

Solution: Added the missing num_steps=cfg.sampler.num_steps parameter to the stochastic sampler configuration.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

No new dependencies required.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Overview

Greptile Summary

Fixes a bug in the CorrDiff generation pipeline where the stochastic sampler was ignoring the configured num_steps parameter and always using the default value of 18 steps.

Changes:

• Added num_steps=cfg.sampler.num_steps parameter to the stochastic_sampler partial() call on line 195
• Brings stochastic sampler configuration in line with deterministic sampler (line 189) which already included this parameter
• Allows users to configure different diffusion step counts for improved performance tuning

Impact:
This is a clean, focused bug fix that resolves a configuration inconsistency. Both sampler functions have num_steps: int = 18 as a default parameter, so the fix ensures the configured value is properly passed through.

Important Files Changed

Filename	Overview
examples/weather/corrdiff/generate.py	Added missing num_steps parameter to stochastic sampler configuration, fixing bug where it always used default 18 steps

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[CharlelieLrt]

LGTM, thanks for the fix!

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

@younes-abid the CI tests are failing because of code-formatting issues. Did you have pre-commit installed as detailed in the contributing guidelines? If not, you need to ruff-format your files

[CharlelieLrt]

@younes-abid following up on this. Could you please format your modified file such that we can then merge your PR? Thanks!

[younes-abid]

ok
I will reformat it