After https://github.com/pytorch/ao/issues/2053 is done, we can remove the old subclass implementation: https://github.com/pytorch/ao/blob/main/torchao/quantization/subclass.py