Update app.py

app.py

@@ -23,7 +23,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("tmberooney/medllama-merged")
 model = AutoModelForCausalLM.from_pretrained("tmberooney/medllama-merged")
-device_map = {"transformer.word_embeddings": "cpu",
+'''device_map = {"transformer.word_embeddings": "cpu",
               "transformer.word_embeddings_layernorm": "cpu",
               "lm_head": "cpu",
               "transformer.h": "cpu",
@@ -34,7 +34,7 @@ model = nn.DataParallel(model)
 for name, param in model.named_parameters():
     if name in device_map:
         param.data = param.to(device=device_map[name])
-#model = model.to('cuda:0')
+#model = model.to('cuda:0')'''