ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update setup.py

Browse source
This commit is contained in:
Daniel Han 2025-07-01 00:34:27 -07:00
commit e73487a4a6
1 changed files with 3 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

3
setup.py
View file

@ -26,6 +26,9 @@ ROOT_DIR = Path(__file__).parent
UNSLOTH_TARGET_DEVICE = os.environ.get("UNSLOTH_TARGET_DEVICE", "cuda")
IS_COLAB = "COLAB_" not in "".join(os.environ.keys())
import torch
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
# Try importing torch
HAS_TORCH = False
HAS_CUDA = True