Delete mem cache handler after training is done

src/otx/algorithms/classification/task.py

@@ -80,6 +80,7 @@
 from otx.api.utils.dataset_utils import add_saliency_maps_to_dataset_item
 from otx.api.utils.labels_utils import get_empty_label
 from otx.cli.utils.multi_gpu import is_multigpu_child_process
+from otx.core.data.caching.mem_cache_handler import MemCacheHandlerSingleton
 
 logger = get_logger()
 RECIPE_TRAIN_TYPE = {
@@ -215,6 +216,8 @@ def train(
 
         results = self._train_model(dataset)
 
+        MemCacheHandlerSingleton.delete()
+
         # Check for stop signal when training has stopped. If should_stop is true, training was cancelled and no new
         if self._should_stop:
             logger.info("Training cancelled.")

src/otx/algorithms/detection/task.py

@@ -65,6 +65,7 @@
 from otx.api.usecases.tasks.interfaces.export_interface import ExportType
 from otx.api.utils.dataset_utils import add_saliency_maps_to_dataset_item
 from otx.cli.utils.multi_gpu import is_multigpu_child_process
+from otx.core.data.caching.mem_cache_handler import MemCacheHandlerSingleton
 
 logger = get_logger()
 
@@ -231,6 +232,8 @@ def train(
         val_dataset.purpose = DatasetPurpose.INFERENCE
         val_preds, val_map = self._infer_model(val_dataset, InferenceParameters(is_evaluation=True))
 
+        MemCacheHandlerSingleton.delete()
+
         preds_val_dataset = val_dataset.with_empty_annotations()
         if self._hyperparams.postprocessing.result_based_confidence_threshold:
             confidence_threshold = 0.0  # Use all predictions to compute best threshold

src/otx/algorithms/segmentation/task.py

@@ -70,6 +70,7 @@
     create_hard_prediction_from_soft_prediction,
 )
 from otx.cli.utils.multi_gpu import is_multigpu_child_process
+from otx.core.data.caching.mem_cache_handler import MemCacheHandlerSingleton
 
 logger = get_logger()
 RECIPE_TRAIN_TYPE = {
@@ -171,6 +172,8 @@ def train(
 
         results = self._train_model(dataset)
 
+        MemCacheHandlerSingleton.delete()
+
         # Check for stop signal when training has stopped. If should_stop is true, training was cancelled and no new
         if self._should_stop:
             logger.info("Training cancelled.")