Skip to content

Commit

Permalink
use Version instead of version.parse
Browse files Browse the repository at this point in the history
  • Loading branch information
qgallouedec committed Nov 20, 2024
1 parent 0c457d0 commit 40baee8
Showing 1 changed file with 1 addition and 2 deletions.
3 changes: 1 addition & 2 deletions src/accelerate/utils/other.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@

import numpy as np
import torch
from packaging import version
from packaging.version import Version
from safetensors.torch import save_file as safe_save_file

Expand Down Expand Up @@ -216,7 +215,7 @@ def save(obj, f, save_on_each_node: bool = False, safe_serialization: bool = Fal

# The following are considered "safe" globals to reconstruct various types of objects when using `weights_only=True`
# These should be added and then removed after loading in the file
np_core = np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core
np_core = np._core if Version(np.__version__) >= Version("2.0.0") else np.core
TORCH_SAFE_GLOBALS = [
# numpy arrays are just numbers, not objects, so we can reconstruct them safely
np_core.multiarray._reconstruct,
Expand Down

0 comments on commit 40baee8

Please sign in to comment.