Fix a bug with a parameter w/o childs in a root module

BloodAxe · Jun 17, 2024 · d73827d · d73827d
1 parent df5ca28
commit d73827d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_toolbelt/optimization/functional.py b/pytorch_toolbelt/optimization/functional.py
@@ -178,7 +178,7 @@ def get_param_group(parameter_name: str, module: nn.Module):
 
     for parameter_name, parameter in parameters:
         total_optimizable_params += parameter.numel()
-        module_name = ".".join(parameter_name.split(".")[:-1])
+        module_name = ".".join(parameter_name.split(".")[:-1]) if "." in parameter_name else parameter_name
         module = recursive_getattr(model, module_name)
 
         param_group: ParametersGroup = get_param_group(parameter_name, module)