This commit is contained in:
Daniel Han 2025-06-30 19:35:59 -07:00
commit a14463004c
2 changed files with 5 additions and 2 deletions

View file

@ -6,9 +6,9 @@ requires = [
"packaging>=24.2",
"setuptools>=77.0.3,<80.0.0",
"setuptools-scm>=8.0",
"torch<=2.7.0",
"wheel",
"jinja2"
"jinja2",
"torch<=2.7.0",
]
build-backend = "setuptools.build_meta"

View file

@ -25,6 +25,7 @@ from setuptools.command.install import install
# This arg is for multi-device
UNSLOTH_TARGET_DEVICE = os.environ.get('UNSLOTH_TARGET_DEVICE', 'cuda')
print("[1] Installing Unsloth...")
def load_module_from_path(module_name, path):
spec = importlib.util.spec_from_file_location(module_name, path)
@ -40,6 +41,8 @@ ROOT_DIR = Path(__file__).parent
# which is not installed yet
ver = load_module_from_path('ver', os.path.join(ROOT_DIR, 'unsloth', 'version.py'))
print("[2] Installing Unsloth...")
def _is_cuda() -> bool:
has_cuda = torch.version.cuda is not None
return UNSLOTH_TARGET_DEVICE == "cuda" and has_cuda