freenode
AI & ML

PyTorch batch-norm stats update overruns undersized tensors

The CPU path indexes running mean and variance by channel count without checking buffer length, causing heap out-of-bounds access.

PyTorch's CPU implementation of batch-normalization statistics updates can read and write past the end of the running mean and variance tensors when those buffers are shorter than the input's channel count.

The operator batch_norm_update_stats uses the input channel dimension to index both running-stat tensors and does not verify that either has enough elements. A minimal case with a four-channel input and one-element running mean and variance triggers a four-byte heap-buffer-overflow under AddressSanitizer on an ASan-enabled CPU build.

Reporter drisspg found the issue during fuzzing and reduced it to a standalone call against the PyTorch project. Batch normalization is widely used in training and inference; frameworks normally keep running-stat shapes aligned with the model, but a mismatched or malformed shape reaches native code that assumes the buffers are large enough. Without a bounds check, the CPU path corrupts adjacent heap memory instead of rejecting the inputs.