Update app.py

app.py

@@ -24,10 +24,10 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 @st.cache(allow_output_mutation=True)
 def load_model():
     model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True)
+    # model = model.to(device)
     tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
-   
-model = model.to(device)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)   
+    
 
 pipe = load_model()