Fix reference to learning_rate

chimezie · Apr 18, 2024 · 6f34e1d · 6f34e1d
1 parent 4af9c08
commit 6f34e1d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_scheduler_config.py b/tests/test_scheduler_config.py
@@ -6,9 +6,8 @@
 
 def load_config(yaml_string, total_iterations):
     config = yaml.load(yaml_string, yaml_loader)
-    learning_rate = config["learning_rate"]
     return SCHEDULE_CONFIGURATION_TYPE_TO_CLASS[
-        config["learning_schedule"]["type"]].from_configuration(learning_rate, config, total_iterations)
+        config["learning_schedule"]["type"]].from_configuration(config.get("learning_rate"), config, total_iterations)
 
 
 class TestCosineWithWarmup: