minor

generate.py

@@ -19,17 +19,14 @@ from utils import StringIteratorIO
 logger = logging.getLogger(__name__)
 
 
-if torch.cuda.is_available():
-    device = "cuda"
-elif torch.backends.mps.is_available():
-    device = "mps"
-else:
-    raise RuntimeError("couldn't find cuda or mps")
-
 logger.warning("Loading model...")
 model_id = "google/gemma-2b-it"
 # model_id = "Qwen/Qwen1.5-0.5B-Chat"
-model = models.transformers(model_id, device=device)
+if torch.backends.mps.is_available():
+    model = models.transformers(model_id, device="mps")
+else:
+    model = models.transformers(model_id, device="cuda")
+
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 sampler = PenalizedMultinomialSampler()
 empty_tokens = [token_id for token_id in range(tokenizer.vocab_size) if not tokenizer.decode([token_id]).strip()]