diff --git a/openfl-workspace/keras_cnn_mnist/src/taskrunner.py b/openfl-workspace/keras_cnn_mnist/src/taskrunner.py
index f07a0c4cfb..1d8d5b7c70 100644
--- a/openfl-workspace/keras_cnn_mnist/src/taskrunner.py
+++ b/openfl-workspace/keras_cnn_mnist/src/taskrunner.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2020-2021 Intel Corporation
+# Copyright (C) 2020-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """You may copy this file as the starting point of your own model."""
diff --git a/openfl-workspace/keras_cnn_with_compression/src/taskrunner.py b/openfl-workspace/keras_cnn_with_compression/src/taskrunner.py
index 60b6c83133..9e9b1eb1c6 100644
--- a/openfl-workspace/keras_cnn_with_compression/src/taskrunner.py
+++ b/openfl-workspace/keras_cnn_with_compression/src/taskrunner.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2020-2021 Intel Corporation
+# Copyright (C) 2020-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """You may copy this file as the starting point of your own model."""
diff --git a/openfl-workspace/keras_nlp/plan/plan.yaml b/openfl-workspace/keras_nlp/plan/plan.yaml
index 5c4f2d3dad..ce7476ab85 100644
--- a/openfl-workspace/keras_nlp/plan/plan.yaml
+++ b/openfl-workspace/keras_nlp/plan/plan.yaml
@@ -14,8 +14,7 @@ collaborator :
   defaults : plan/defaults/collaborator.yaml
   template : openfl.component.Collaborator
   settings :
-    epochs_per_round : 10
-    polling_interval : 4
+    db_store_rounds: 2
     delta_updates    : false
     opt_treatment    : RESET
 
diff --git a/openfl-workspace/keras_nlp/src/dataloader.py b/openfl-workspace/keras_nlp/src/dataloader.py
index bbda8a193e..d7180b4d3c 100644
--- a/openfl-workspace/keras_nlp/src/dataloader.py
+++ b/openfl-workspace/keras_nlp/src/dataloader.py
@@ -6,6 +6,7 @@
 """
 from logging import getLogger
 from typing import Optional
+from typing import Iterator
 from typing import Tuple
 from typing import Union
 
@@ -55,7 +56,7 @@ def get_feature_shape(self) -> Tuple[int, ...]:
         """Get the shape of an example feature array."""
         return self.X_train[0].shape
 
-    def get_train_loader(self, batch_size: Optional[int] = None):
+    def get_train_loader(self, batch_size: Optional[int] = None) -> Iterator[Tuple[np.ndarray]]:
         """
         Get training data loader.
 
@@ -66,7 +67,7 @@ def get_train_loader(self, batch_size: Optional[int] = None):
         return self._get_batch_generator(X1=self.X_train[0], X2=self.X_train[1],
                                          y=self.y_train, batch_size=batch_size)
 
-    def get_valid_loader(self, batch_size: Optional[int] = None):
+    def get_valid_loader(self, batch_size: Optional[int] = None) -> Iterator[Tuple[np.ndarray]]:
         """
         Get validation data loader.
 
@@ -98,7 +99,7 @@ def get_valid_data_size(self) -> int:
     def _batch_generator(X1: np.ndarray, X2: np.ndarray,
                          y: np.ndarray, idxs: np.ndarray,
                          batch_size: int,
-                         num_batches: int):
+                         num_batches: int) -> Iterator[Tuple[np.ndarray]]:
         """
         Generate batch of data.
 
diff --git a/openfl-workspace/keras_nlp/src/taskrunner.py b/openfl-workspace/keras_nlp/src/taskrunner.py
index 6e2d3a8c8a..4463dcebe8 100644
--- a/openfl-workspace/keras_nlp/src/taskrunner.py
+++ b/openfl-workspace/keras_nlp/src/taskrunner.py
@@ -1,4 +1,4 @@
-"""Copyright (C) 2020-2021 Intel Corporation
+"""Copyright (C) 2020-2024 Intel Corporation
    SPDX-License-Identifier: Apache-2.0
 
 Licensed subject to the terms of the separately executed evaluation