diff --git a/embedlab.py b/embedlab.py index 92d2b11..9db35ba 100755 --- a/embedlab.py +++ b/embedlab.py @@ -31,7 +31,12 @@ LAST_MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cac HISTORY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "history") COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"] -device = "cuda" if torch.cuda.is_available() else "cpu" +if torch.cuda.is_available(): + device = "cuda" +elif torch.backends.mps.is_available(): # Apple Silicon GPU + device = "mps" +else: + device = "cpu" _loaded = {} # loaded models: name -> (tokenizer, model) _params = {name: dict(d) for name, d in MODELS.items()} # live, per-model parameters current = None # currently selected model