fix: error in raytune search space (#170)

Former-commit-id: 359b892
jpata · Feb 15, 2023 · 6ca85cb · 6ca85cb
1 parent 5762af0
commit 6ca85cb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlpf/raytune/search_space.py b/mlpf/raytune/search_space.py
@@ -209,7 +209,7 @@ def set_raytune_search_parameters(search_space, config):
         config["parameters"]["output_decoding"]["phi_hidden_dim"] = search_space["out_hidden_dim"]
         config["parameters"]["output_decoding"]["energy_hidden_dim"] = search_space["out_hidden_dim"]
 
-    if "hidden_dim" in search_space.keys():
+    if "out_num_layers" in search_space.keys():
         config["parameters"]["output_decoding"]["id_num_layers"] = search_space["out_num_layers"]
         config["parameters"]["output_decoding"]["charge_num_layers"] = search_space["out_num_layers"]
         config["parameters"]["output_decoding"]["pt_num_layers"] = search_space["out_num_layers"]