changes class_implementation to init_class in gpxq_mode

src/brevitas/graph/gpfq.py

@@ -59,8 +59,7 @@ def __init__(
 
         self.orig_forward = self.model.forward
         self.model.forward = self.catch_stopfwd
-        self.class_implementation = GPFQ
-        GPFQ.p = p
+        self.p = p
 
     def catch_stopfwd(self, *args, **kwargs):
         # Collect quant input
@@ -95,14 +94,29 @@ def catch_stopfwd(self, *args, **kwargs):
                 gpxq_class.disable_pre_forward_hook = False
             return out
 
+    def init_class(self, layer, name, act_order, parallel_layers, create_weight_orig):
+        return GPFQ(
+            layer=layer,
+            name=name,
+            act_order=act_order,
+            parallel_layers=parallel_layers,
+            create_weight_orig=create_weight_orig,
+            p=self.p)
+
 
 class GPFQ(GPxQ):
     """
     Based on https://github.com/YixuanSeanZhou/Quantized_Neural_Nets/tree/main
     """
-    p = 0.25
 
-    def __init__(self, layer, name, act_order, parallel_layers=1, create_weight_orig=True) -> None:
+    def __init__(
+            self,
+            layer,
+            name,
+            act_order,
+            parallel_layers=1,
+            create_weight_orig=True,
+            p=0.25) -> None:
 
         if act_order:
             raise ValueError("Act_order is not supported in GPFQ")
@@ -111,7 +125,7 @@ def __init__(self, layer, name, act_order, parallel_layers=1, create_weight_orig
         self.float_input = None
         self.quantized_input = None
         self.index_computed = False
-        self.p = GPFQ.p
+        self.p = p
 
     def update_batch(self, module, input, current_layer):
         if self.disable_pre_forward_hook:

src/brevitas/graph/gptq.py

@@ -67,8 +67,6 @@ def __init__(
         self.model.forward = self.catch_stopfwd
         # How many subblock to use during GPTQ for each layer
         self.num_blocks = num_blocks
-        self.class_implementation = GPTQ
-        GPTQ.num_blocks = num_blocks
 
     def catch_stopfwd(self, *args, **kwargs):
         try:
@@ -85,6 +83,15 @@ def catch_stopfwd(self, *args, **kwargs):
                     gpxq_class.disable_pre_forward_hook = False
                 return out
 
+    def init_class(self, layer, name, act_order, parallel_layers, create_weight_orig):
+        return GPTQ(
+            layer=layer,
+            name=name,
+            act_order=act_order,
+            parallel_layers=parallel_layers,
+            create_weight_orig=create_weight_orig,
+            num_blocks=self.num_blocks)
+
 
 class GPTQ(GPxQ):
     """
@@ -104,15 +111,22 @@ class GPTQ(GPxQ):
     See the License for the specific language governing permissions and
     limitations under the License.
     """
-    num_blocks = 100
 
-    def __init__(self, layer, name, act_order, parallel_layers=1, create_weight_orig=True) -> None:
+    def __init__(
+            self,
+            layer,
+            name,
+            act_order,
+            parallel_layers=1,
+            create_weight_orig=True,
+            num_blocks=100) -> None:
         super().__init__(layer, name, act_order, parallel_layers, create_weight_orig)
 
         dev = self.layer.weight.device
 
         # Define how many columns to update in each mini-block
-        self.blocksize = math.ceil(self.columns / GPTQ.num_blocks)
+        self.blocksize = math.ceil(self.columns / num_blocks)
+        self.num_blocks = num_blocks
 
         # Initialize Hessian matrix and counter. We need it in float32 to compute the inverse
         self.H = torch.zeros((self.groups, self.columns, self.columns),

src/brevitas/graph/gpxq.py

@@ -98,7 +98,7 @@ def __enter__(self):
 
                 # Attach hooks for GPTQ
                 if self._is_module_supported(module):
-                    gpxq = self.class_implementation(
+                    gpxq = self.init_class(
                         module,
                         name,
                         act_order=self.act_order,