Skip to content

Commit

Permalink
[hotfix] improve compatibility (hpcaitech#6165)
Browse files Browse the repository at this point in the history
  • Loading branch information
ver217 authored Dec 23, 2024
1 parent 5f82bfa commit 8b0ed61
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions colossalai/utils/safetensors.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214
import json
import warnings
from dataclasses import asdict, dataclass
from typing import Dict, List, Optional, Tuple

Expand All @@ -8,8 +9,10 @@

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

Expand Down

0 comments on commit 8b0ed61

Please sign in to comment.