diff --git a/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldProcessingService.java b/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldProcessingService.java
index 9c9020686b..1ab8f5a076 100644
--- a/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldProcessingService.java
+++ b/application/src/main/java/org/thingsboard/server/service/cf/DefaultCalculatedFieldProcessingService.java
@@ -106,10 +106,8 @@ public class DefaultCalculatedFieldProcessingService extends AbstractCalculatedF
         if (ctx.getCfType() != CalculatedFieldType.PROPAGATION) {
             return Optional.empty();
         }
-        PropagationArgumentEntry argumentEntry = (PropagationArgumentEntry)
-                resolveArgumentValue(PROPAGATION_CONFIG_ARGUMENT, fetchPropagationCalculatedFieldArgument(ctx, entityId));
-        argumentEntry.setPartitionStateRestore(true);
-        return Optional.of(argumentEntry);
+        return Optional.of((PropagationArgumentEntry)
+                resolveArgumentValue(PROPAGATION_CONFIG_ARGUMENT, fetchPropagationCalculatedFieldArgument(ctx, entityId)));
     }
 
     @Override
diff --git a/application/src/main/java/org/thingsboard/server/service/cf/ctx/state/propagation/PropagationCalculatedFieldState.java b/application/src/main/java/org/thingsboard/server/service/cf/ctx/state/propagation/PropagationCalculatedFieldState.java
index 3fe06ea4c9..0f1c44d0e2 100644
--- a/application/src/main/java/org/thingsboard/server/service/cf/ctx/state/propagation/PropagationCalculatedFieldState.java
+++ b/application/src/main/java/org/thingsboard/server/service/cf/ctx/state/propagation/PropagationCalculatedFieldState.java
@@ -69,6 +69,7 @@ public class PropagationCalculatedFieldState extends ScriptCalculatedFieldState
         super.init(restored);
         if (restored) {
             cfProcessingService.fetchPropagationArgumentFromDb(ctx, entityId).ifPresent(fromDb -> {
+                fromDb.setPartitionStateRestore(true);
                 var update = update(Map.of(PROPAGATION_CONFIG_ARGUMENT, fromDb), ctx);
                 if (update.isEmpty()) {
                     return;