diff --git a/unsloth/_auto_install.py b/unsloth/_auto_install.py index 27b23ed476..e8afc32541 100644 --- a/unsloth/_auto_install.py +++ b/unsloth/_auto_install.py @@ -15,12 +15,12 @@ try: import torch except: raise ImportError('Install torch via `pip install torch`') from packaging.version import Version as V -v = V(torch.__version__) +import re +v = V(re.match(r"[0-9\.]{3,}", torch.__version__).group(0)) cuda = str(torch.version.cuda) is_ampere = torch.cuda.get_device_capability()[0] >= 8 USE_ABI = torch._C._GLIBCXX_USE_CXX11_ABI -if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8"): - raise RuntimeError(f"CUDA = {cuda} not supported!") +if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!") if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!") elif v <= V('2.1.1'): x = 'cu{}{}-torch211' elif v <= V('2.1.2'): x = 'cu{}{}-torch212' @@ -34,7 +34,6 @@ elif v < V('2.7.9'): x = 'cu{}{}-torch270' elif v < V('2.8.0'): x = 'cu{}{}-torch271' elif v < V('2.8.9'): x = 'cu{}{}-torch280' else: raise RuntimeError(f"Torch = {v} too new!") -if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8"): - raise RuntimeError(f"CUDA = {cuda} not supported!") +if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!") x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "") print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"') \ No newline at end of file