Skip to content

Commit 8b0ed61

Browse files
authored
[hotfix] improve compatibility (#6165)
1 parent 5f82bfa commit 8b0ed61

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

colossalai/utils/safetensors.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214
22
import json
3+
import warnings
34
from dataclasses import asdict, dataclass
45
from typing import Dict, List, Optional, Tuple
56

@@ -8,8 +9,10 @@
89

910
try:
1011
from tensornvme.async_file_io import AsyncFileWriter
11-
except ModuleNotFoundError:
12-
raise ModuleNotFoundError("Please install tensornvme to use NVMeOptimizer")
12+
except Exception:
13+
warnings.warn(
14+
"Please install the latest tensornvme to use async save. pip install git+https://github.com/hpcaitech/TensorNVMe.git"
15+
)
1316
_TYPES_INV = {v: k for k, v in _TYPES.items()}
1417
import io
1518

0 commit comments

Comments
 (0)