[ao] Fixing tests for block pruning shapes (#87326)
The current unittests were only checking the tensors whose shapes were already multiples of the block size. That caused some hidden bugs to creep in. Specifically, for the shapes that would require padding for the mask/data, the sparsifier would try to apply shape-mismatching tensors onto each other. This caused segfaults as well as silent failures.
This makes minor adjustments to the code to make sure the masks and data shapes are aligned, as well as fixing the tests to catch this.
Test Plan:
```python
python test/test_ao_sparsity.py
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87326
Approved by: https://github.com/jcaip
