diff --git a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py
index 98aea2be5..f25e2f5fa 100644
--- a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py
+++ b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py
@@ -26,6 +26,7 @@
 import dataclasses
 import functools
 import json
+import os
 import threading
 import time
 from typing import Any, List, Mapping, Optional, Tuple, Union