Update app.py

app.py

@@ -7,10 +7,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def load_model():
-    model_id = "nanochat-students/chat-d20"
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False)
-    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=False, dtype=torch.bfloat16).to(device)
+    model_id="karpathy/nanochat-d32"
+    revision="refs/pr/1"
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False, revision=revision)
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=False, dtype=torch.bfloat16, revision=revision).to(device)
+
     model.eval()
 
     return tokenizer, model