From 8b0ed614909154388bdcb3884324a159c162f524 Mon Sep 17 00:00:00 2001 From: Hongxin Liu Date: Mon, 23 Dec 2024 18:57:08 +0800 Subject: [PATCH] [hotfix] improve compatibility (#6165) --- colossalai/utils/safetensors.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/colossalai/utils/safetensors.py b/colossalai/utils/safetensors.py index 8ce6d7335..cf404038c 100644 --- a/colossalai/utils/safetensors.py +++ b/colossalai/utils/safetensors.py @@ -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 @@ -8,8 +9,10 @@ from safetensors.torch import _TYPES, load_file, safe_open 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