fix issues in configure_module_default_qconfigs (#359)

Author: bfineran

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

@@ -275,10 +275,12 @@ def configure_module_qat_wrappers(module: Module):
 
     :param module: module to potentially wrap the submodules of
     """
-    for submodule in module.modules():
-        for child_name, child_module in module.named_children():
-            if hasattr(child_module, "wrap_qat") and child_module.wrap_qat:
-                setattr(submodule, child_name, QATWrapper.from_module(child_module))
+    # wrap any children of the given module as a QATWrapper if required
+    for child_name, child_module in module.named_children():
+        if hasattr(child_module, "wrap_qat") and child_module.wrap_qat:
+            setattr(module, child_name, QATWrapper.from_module(child_module))
+        # recurse on child module
+        configure_module_qat_wrappers(child_module)
 
 
 def configure_module_default_qconfigs(module: Module):

tests/sparseml/pytorch/utils/quantization/test_helpers.py

@@ -82,13 +82,13 @@ def _count_submodule_instances(module, clazz):
     reason="torch quantization not available",
 )
 def test_configure_module_qat_wrappers():
-    module = _ModuleWrapper(_QATMatMul())
+    module = _ModuleWrapper(_ModuleWrapper(_QATMatMul()))
 
     assert not _module_has_quant_stubs(module)
 
     configure_module_qat_wrappers(module)
 
-    qat_matmul = module.module
+    qat_matmul = module.module.module
 
     assert isinstance(qat_matmul, QATWrapper)
     assert _module_has_quant_stubs(module)