From cf7448877d53908cb1949a3f9c179039902e0c2a Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Fri, 1 Mar 2024 10:44:40 +0800
Subject: [PATCH] Revert "remove `self._set_prop`"

This reverts commit 3be3f85119d4180eef88d8945f3d1554129ac9aa.
---
 bundle/python_bundle_workflow/scripts/inference.py | 6 +++++-
 bundle/python_bundle_workflow/scripts/train.py     | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/bundle/python_bundle_workflow/scripts/inference.py b/bundle/python_bundle_workflow/scripts/inference.py
index 19f9aa21ee..06ce1aff5f 100644
--- a/bundle/python_bundle_workflow/scripts/inference.py
+++ b/bundle/python_bundle_workflow/scripts/inference.py
@@ -63,6 +63,7 @@ def __init__(self, dataset_dir: str = "."):
             nib.save(n, os.path.join(dataset_dir, f"seg{i:d}.nii.gz"))
 
         self._props = {}
+        self._set_props = {}
         self.dataset_dir = dataset_dir
 
     def initialize(self):
@@ -75,7 +76,7 @@ def finalize(self):
         pass
 
     def _set_property(self, name, property, value):
-        self._props[name] = value
+        self._set_props[name] = value
 
     def _get_property(self, name, property):
         """
@@ -89,6 +90,9 @@ def _get_property(self, name, property):
         value = None
         if name in self._props:
             value = self._props[name]
+        elif name in self._set_props:
+            value = self._set_props[name]
+            self._props[name] = value
         else:
             try:
                 value = getattr(self, f"get_{name}")()
diff --git a/bundle/python_bundle_workflow/scripts/train.py b/bundle/python_bundle_workflow/scripts/train.py
index bc8a47afa6..16fa85b324 100644
--- a/bundle/python_bundle_workflow/scripts/train.py
+++ b/bundle/python_bundle_workflow/scripts/train.py
@@ -72,6 +72,7 @@ def __init__(self, dataset_dir: str = "."):
 
         # define buckets to store the generated properties and set properties
         self._props = {}
+        self._set_props = {}
         self.dataset_dir = dataset_dir
 
         # besides the predefined properties, this bundle workflow can also provide `network`, `loss`, `optimizer`
@@ -89,7 +90,7 @@ def finalize(self):
         set_determinism(None)
 
     def _set_property(self, name, property, value):
-        self._props[name] = value
+        self._set_props[name] = value
 
     def _get_property(self, name, property):
         """
@@ -103,6 +104,9 @@ def _get_property(self, name, property):
         value = None
         if name in self._props:
             value = self._props[name]
+        elif name in self._set_props:
+            value = self._set_props[name]
+            self._props[name] = value
         else:
             try:
                 value = getattr(self, f"get_{name}")()