Repository: pytorch
Issue: [Operator Benchmark] [CPU]OP: weight_norm_sparsifier_step functionality failed with parts of long config shapes
Labels: oncall: quantization

### 🐛 Describe the bug

OP: weight_norm_sparsifier_step functionality failed with parts of long config shapes on CPU.

Reproduce:
```
pip3 install torch
```
Please refer to https://github.com/DiweiSun/pytorch/tree/molly/op_microbenchmark/benchmarks/operator_benchmark to install pt extension for operator benchmark.

And then run cmd:
```
python -m pt.ao_sparsifier_test  --tag_filter long
```


Error logs:
```
# Benchmarking PyTorch: weight_norm_sparsifier_step
# Mode: Eager
# Name: weight_norm_sparsifier_step_M(255,324)_SL1.0_SBS(1,4)_ZPBNone
# Input: M: (255, 324), SL: 1.0, SBS: (1, 4), ZPB: None
Forward Execution Time (us) : 31.321

# Benchmarking PyTorch: weight_norm_sparsifier_step
# Mode: Eager
# Name: weight_norm_sparsifier_step_M(255,324)_SL1.0_SBS(8,1)_ZPB0
# Input: M: (255, 324), SL: 1.0, SBS: (8, 1), ZPB: 0
Forward Execution Time (us) : 30.981

# Benchmarking PyTorch: weight_norm_sparsifier_step
Traceback (most recent call last):
  File "/home/pytorch/miniforge/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/pytorch/miniforge/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/pt/ao_sparsifier_test.py", line 54, in <module>
    op_bench.benchmark_runner.main()
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/benchmark_runner.py", line 179, in main
    benchmark_core.BenchmarkRunner(args).run()
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/benchmark_core.py", line 467, in run
    launch_func(
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/benchmark_core.py", line 280, in _launch_forward
    forward_time = timeit.timeit(
  File "/home/pytorch/miniforge/lib/python3.10/timeit.py", line 234, in timeit
    return Timer(stmt, setup, timer, globals).timeit(number)
  File "/home/pytorch/miniforge/lib/python3.10/timeit.py", line 178, in timeit
    timing = self.inner(it, self.timer)
  File "<timeit-src>", line 6, in inner
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/benchmark_pytorch.py", line 160, in run_forward
    self.output = self.op_bench.forward_impl()
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/benchmark_pytorch.py", line 66, in forward_impl
    return self.forward(*self.get_inputs())
  File "/home/pytorch/molly/04-pytorch/benchmarks/operator_benchmark/pt/ao_sparsifier_test.py", line 46, in forward
    self.sparsifier.step()
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/ao/pruning/sparsifier/base_sparsifier.py", line 350, in step
    self.update_mask(**config)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py", line 193, in update_mask
    ww = self.norm_fn(getattr(module, tensor_name))
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/utils/parametrize.py", line 379, in get_parametrized
    return parametrization()
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/utils/parametrize.py", line 276, in forward
    x = self[0](self.original)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/pytorch/miniforge/lib/python3.10/site-packages/torch/ao/pruning/sparsifier/utils.py", line 130, in forward
    assert self.mask.shape == x.shape
AssertionError
```


### Versions

[pip3] flake8==6.1.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.15.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] mypy==1.10.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.0
[pip3] optree==0.12.1
[pip3] torch==2.4.0
[pip3] triton==3.0.0
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] optree                    0.12.1                   pypi_0    pypi
[conda] torch                     2.4.0                    pypi_0    pypi
[conda] torchfix                  0.4.0                    pypi_0    pypi
[conda] triton                    3.0.0                    pypi_0    pypi

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim