[BUGFIX][BACKPORT]fixing breaking change introduced in #17123 when batch_axis=0

[access2rohit]

Description

Backport #19267

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[mxnet-bot]

Hey @access2rohit , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [sanity, edge, centos-cpu, unix-gpu, unix-cpu, website, centos-gpu, miscellaneous, clang, windows-cpu, windows-gpu]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[sandeep-krishnamurthy]

@mxnet-bot run ci [unix-cpu]

[sandeep-krishnamurthy]

@samskalicky

[mxnet-bot]

Jenkins CI successfully triggered : [unix-cpu]

[access2rohit]

@samskalicky can you merge ?