diff --git a/python/ray/air/_internal/checkpoint_manager.py b/python/ray/air/_internal/checkpoint_manager.py
index 52dda666f9c9..165efc790c3c 100644
--- a/python/ray/air/_internal/checkpoint_manager.py
+++ b/python/ray/air/_internal/checkpoint_manager.py
@@ -10,10 +10,10 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import ray
+from ray._private.dict import flatten_dict
 from ray.air import Checkpoint, CheckpointConfig
 from ray.air.config import MAX
 from ray.air._internal.util import is_nan
-from ray.tune.utils.util import flatten_dict
 from ray.util import log_once