Add SAWB-based recomputation of weights and refactor whole process

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/quant/quantizers.py

@@ -428,7 +428,7 @@ def set_quantizer(self):
             else:
                 self.quantizer = SAWBSTE
 
-    def forward(self, input_tensor):
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         input_tensor = self.quantizer.apply(
             input_tensor,
             self.num_bits,