Add weight qconfig properties to forward_fn in case it contains weights (#940)

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

Author: anmarques

src/sparseml/pytorch/sparsification/quantization/helpers.py

@@ -256,6 +256,13 @@ def __init__(
         num_input_quant_stubs = num_inputs + len(self.kwarg_input_names)
 
         self.forward_fn = forward_fn
+        # Add weight qconfig to forward_fn (in case it has weights)
+        qconfig_ = get_qat_qconfig(qproperties)
+        qconfig = torch_quantization.QConfig(
+            activation=torch.nn.Identity,
+            weight=qconfig_.weight,
+        )
+        self.forward_fn.qconfig = qconfig
 
         self.input_qconfigs = self._load_qconfigs(
             name="input_qconfigs",