Tensor size mismatch when using static quantization

[yiliu30]

At static_quant, if we change the ToyLinearModel(1024, 1024, 1024) to ToyLinearModel(), it causes the below issue:

Traceback (most recent call last):
  File "projs/torchao/tutorials/calibration_flow/static_quant.py", line 132, in <module>
    m(*example_inputs)
  File "conda_env/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "projs/torchao/tutorials/calibration_flow/static_quant.py", line 109, in forward
    x = self.linear2(x)
        ^^^^^^^^^^^^^^^
  File "conda_env/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "projs/torchao/tutorials/calibration_flow/static_quant.py", line 27, in forward
    observed_weight = self.weight_obs(self.weight)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/ao/quantization/observer.py", line 739, in forward
    return self._forward(x_orig)
           ^^^^^^^^^^^^^^^^^^^^^
  File "conda_env/torch/ao/quantization/observer.py", line 762, in _forward
    min_val = torch.min(min_val_cur, min_val)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The size of tensor a (64) must match the size of tensor b (32) at non-singleton dimension 0

cc @jerryzh168

[jerryzh168]

thanks, there is a bug in replacement_fn I think, fixed in #650

[jerryzh168]

should be fixed now