fix how config's optimizer creates optimizer

GFNOrg · Apr 5, 2023 · a6ffcdf · a6ffcdf
1 parent 5b15d01
commit a6ffcdf
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/scripts/configs/optim.py b/scripts/configs/optim.py
@@ -25,14 +25,21 @@ def get_params(
                 "params": [
                     val
                     for key, val in parametrization.parameters.items()
-                    if key != "logZ"
+                    if "logZ" not in key
                 ],
                 "lr": self.lr,
             }
         ]
         if "logZ" in parametrization.parameters:
             params.append(
-                {"params": [parametrization.parameters["logZ"]], "lr": self.lr_Z}
+                {
+                    "params": [
+                        val
+                        for key, val in parametrization.parameters.items()
+                        if "logZ" in key
+                    ],
+                    "lr": self.lr_Z,
+                }
             )
         else:
             self.lr_Z = None