diff --git a/test/test_modules/test_training.py b/test/test_modules/test_training.py
index 66ba0709c21b105bd798cd35f20715e6c0a83177..eb5dfe5adb170981d5d67c94ca1fbcb55e326550 100644
--- a/test/test_modules/test_training.py
+++ b/test/test_modules/test_training.py
@@ -156,7 +156,7 @@ class TestTraining:
         obj.data_store.set("model_path", model_path, "general")
         obj.data_store.set("model_name", os.path.join(model_path, "test_model.h5"), "general.model")
         obj.data_store.set("batch_size", 256, "general")
-        obj.data_store.set("epochs", 2, "general.model")
+        obj.data_store.set("epochs", 2, "general")
         clbk, hist, lr = callbacks
         obj.data_store.set("callbacks", clbk, "general.model")
         obj.data_store.set("lr_decay", lr, "general.model")