Segfault in bincount

[SleepyMug]

🐛 Describe the bug

Function torch.bincount contains segmentation fault.

Example to reproduce

import torch

self = torch.full([3], 2550, dtype=torch.int64, requires_grad=False)
weights = torch.full([3, 1, 3, 0, 0, 0, 1, 1], -4620, dtype=torch.int64, requires_grad=False)
minlength = 9711
torch.bincount(self, weights, minlength)

Result

Segmentation fault

Expected Behavior

Throwing a Python Exception

Notes

This bug was discovered using Atheris.

Versions

Is debug build: N/A
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: N/A
OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: 11.0.1-2
CMake version: version 3.18.4
Libc version: glibc-2.31
Python version: 3.9.7 (default, Sep 16 2021, 13:09:58) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.0-11-amd64-x86_64-with-glibc2.31
Is CUDA available: N/A
CUDA runtime version: Could not collect
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: N/A
Versions of relevant libraries:
[pip3] numpy==1.22.3
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchvision==0.12.0
[conda] blas 1.0 mkl
[conda] cpuonly 2.0 0 pytorch
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py39h7f8727e_0
[conda] mkl_fft 1.3.1 py39hd3c417c_0
[conda] mkl_random 1.2.2 py39h51133e4_0
[conda] numpy 1.21.2 py39h20f2e39_0
[conda] numpy-base 1.21.2 py39h79a1101_0
[conda] pytorch 1.11.0 py3.9_cpu_0 pytorch
[conda] pytorch-mutex 1.0 cpu pytorch
[conda] torchaudio 0.11.0 py39_cpu pytorch
[conda] torchvision 0.12.0 py39_cpu pytorch

[ajedd077]

This is caused by the zero in weights sizes, making the number of elements set to zero, but weights.size(0) == self.size(0). A possible solution is to add a numel check to weights in Tensor _bincount_cpu_template.
From:

  bool has_weights = weights.defined();
  if (has_weights && weights.size(0) != self.size(0)) {
    AT_ERROR("input and weights should have the same length");
  }

to :

  bool has_weights = weights.defined();
  if (has_weights && weights.size(0) != self.size(0)  && weights.numel() != 0) {
    AT_ERROR("input and weights should have the same length");
  }

Excuse my ignorance as it is the first-ever open-source comment. I know it is trivial - but I would like to contribute frequently to this project. Any guidelines or suggestions are appreciated.