Fix wrong detection of depthwise conv on neon (#55794)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/54136
tldr: dephwise conv require that the nb of output channel is 1.
The code here only handles this case and previously, all but the first output channel were containing uninitialized memory. The nans from the issue were random due to the allocation of a torch.empty() that was sometimes returning non-nan memory.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55794
Reviewed By: ngimel
Differential Revision: D27711717
Pulled By: albanD
fbshipit-source-id: 00eac3fd59db1d09fe7bab89427b105a019e7a5d
