diff --git a/docs/doc_usage_cmd.md b/docs/doc_usage_cmd.md
index ee814fdd3..28325b5af 100644
--- a/docs/doc_usage_cmd.md
+++ b/docs/doc_usage_cmd.md
@@ -43,9 +43,11 @@ To run DomainLab, the minimum necessary parameters are:
       ```yaml
       gamma_reg:
         dann: 0.05
-        diva: 0.2
+        dial: 0.2
         default: 0.1 # value for every other instance
       ```  
+Gamma reg is available for the trainers, as well as the 
+dann and jigen model. 
 
 - **Early Stopping (`--es`):** Steps for early stopping.
 - **Random Seed (`--seed`):** Seed for reproducibility.
diff --git a/domainlab/arg_parser.py b/domainlab/arg_parser.py
index 32faa48be..1b386c698 100644
--- a/domainlab/arg_parser.py
+++ b/domainlab/arg_parser.py
@@ -15,8 +15,22 @@
 class ParseValuesOrKeyValuePairs(argparse.Action):
     """Class used for arg parsing where values are provided in a key value format"""
 
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(self, parser: argparse.ArgumentParser,
+    namespace: argparse.Namespace, values: str, option_string: str = None):
+        """
+            Handle parsing of key value pairs, or a single value instead
+
+            Args:
+                parser (argparse.ArgumentParser): The ArgumentParser object.
+                namespace (argparse.Namespace): The namespace object to store parsed values.
+                values (str): The string containing key=value pairs or a single float value.
+                option_string (str, optional): The option string that triggered this action (unused).
+
+            Raises:
+                ValueError: If the values cannot be parsed to float.
+        """
         if "=" in values:
+            my_dict = {}        if "=" in values:
             my_dict = {}
             for kv in values.split(","):
                 k, v = kv.split("=")