Fix for handling batch size 0. (#34599)
Summary:
Separating this out in a different diff, however since most of the
xnnpack integration is not tested until the PR https://github.com/pytorch/pytorch/issues/34047, this was not
caught till then.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/34599
Test Plan: Tested in test/test_xnnpack_integration.py via https://github.com/pytorch/pytorch/issues/34047.
Differential Revision: D20391000
Pulled By: kimishpatel
fbshipit-source-id: 596a3e54445072ab63f700d425d07c7f44586683