diff --git a/setup.py b/setup.py index fc4c113166..b411d4d18c 100644 --- a/setup.py +++ b/setup.py @@ -30,9 +30,7 @@ IS_COLAB = "COLAB_" not in "".join(os.environ.keys()) HAS_TORCH = False HAS_CUDA = True HAS_HIP = False -CUDA_HOME = None -ROCM_HOME = None -torch = None +torch = None if importlib.util.find_spec("torch") is not None: try: import torch @@ -44,9 +42,13 @@ if importlib.util.find_spec("torch") is not None: except Exception as e: print(f"Unsloth: Importing torch failed with error = {str(e)}.") HAS_TORCH = False + CUDA_HOME = None + ROCM_HOME = None +else: + CUDA_HOME = None + ROCM_HOME = None pass - def load_module_from_path(module_name, path): spec = importlib.util.spec_from_file_location(module_name, path) module = importlib.util.module_from_spec(spec)