From b0f34776dfae2a24b4b826d4b064d5865355505c Mon Sep 17 00:00:00 2001 From: digger yu Date: Wed, 28 Jun 2023 16:17:13 +0800 Subject: [PATCH] fix CheckpointIndexFile is not defined --- colossalai/checkpoint_io/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py index 3dada00cd9b5..f28b3e82cc9e 100644 --- a/colossalai/checkpoint_io/utils.py +++ b/colossalai/checkpoint_io/utils.py @@ -11,6 +11,7 @@ from torch.optim import Optimizer from colossalai.tensor.d_tensor.d_tensor import DTensor +from .index_file import CheckpointIndexFile SAFE_WEIGHTS_NAME = "model.safetensors" WEIGHTS_NAME = "pytorch_model.bin"