[AutoDeploy[bug-9379] Fixing sharding-related bug.

[greg-kwasniewski1]

🚀 The feature, motivation and pitch

See bug description:
https://nvbugspro.nvidia.com/bug/5680312
#9379

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.