Skip to content

Commit

Permalink
Drop torch re-imports in npu and mlu paths (#2856)
Browse files Browse the repository at this point in the history
Signed-off-by: Dmitry Rogozhkin <dmitry.v.rogozhkin@intel.com>
  • Loading branch information
dvrogozh authored Jun 14, 2024
1 parent 3b5a00e commit 13ca7dc
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions src/accelerate/utils/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,6 @@ def is_mlu_available(check_device=False):
if importlib.util.find_spec("torch_mlu") is None:
return False

import torch
import torch_mlu # noqa: F401

if check_device:
Expand All @@ -359,10 +358,9 @@ def is_mlu_available(check_device=False):
@lru_cache
def is_npu_available(check_device=False):
"Checks if `torch_npu` is installed and potentially if a NPU is in the environment"
if importlib.util.find_spec("torch") is None or importlib.util.find_spec("torch_npu") is None:
if importlib.util.find_spec("torch_npu") is None:
return False

import torch
import torch_npu # noqa: F401

if check_device:
Expand Down

0 comments on commit 13ca7dc

Please sign in to comment.