osx

src/autotrain/cli/run_llm.py

@@ -407,10 +407,16 @@ def __init__(self, args):
                     break
                 print(f"Bot: {tgi.chat(prompt)}")
 
-        if not torch.cuda.is_available():
-            raise ValueError("No GPU found. Please install CUDA and try again.")
+        cuda_available = torch.cuda.is_available()
+        mps_available = torch.mps.is_available()
 
-        self.num_gpus = torch.cuda.device_count()
+        if not cuda_available and not mps_available:
+            raise ValueError("No GPU/MPS device found. LLM training requires an accelerator")
+
+        if cuda_available:
+            self.num_gpus = torch.cuda.device_count()
+        elif mps_available:
+            self.num_gpus = 1
 
     def run(self):
         from autotrain.backend import EndpointsRunner, SpaceRunner

src/autotrain/trainers/clm/__main__.py

@@ -129,6 +129,7 @@ def train(config):
             trust_remote_code=True,
             use_flash_attention_2=config.use_flash_attention_2,
         )
+
     else:
         model = AutoModelForCausalLM.from_pretrained(
             config.model,