diff --git a/test/test_run_modules/test_training.py b/test/test_run_modules/test_training.py
index 51ea1cd344c1ff1899af818c6b38a2cbb93b733a..b16c0c2586f87af8368ac0059edc8a3997780f69 100644
--- a/test/test_run_modules/test_training.py
+++ b/test/test_run_modules/test_training.py
@@ -194,6 +194,7 @@ class TestTraining:
         obj.data_store.set("data_collection", data_collection, "general.val")
         obj.data_store.set("data_collection", data_collection, "general.test")
         obj.model.compile(**obj.model.compile_options)
+        keras.utils.get_custom_objects().update(obj.model.custom_objects)
         return obj
 
     @pytest.fixture