Update app.py

app.py

@@ -10,7 +10,7 @@ print(device)
 
 model_name = "bigscience/bloom-1b"
 
-model = BloomForCausalLM.from_pretrained(model_name,device_map="auto", load_in_8bit=True)
+model = BloomForCausalLM.from_pretrained(model_name).to(device)
 tokenizer = BloomTokenizerFast.from_pretrained(model_name)
 
 # Define the pipeline for Gradio purpose
@@ -19,7 +19,7 @@ def beam_gradio_pipeline(prompt,length=500):
 
    result_length = length
 
-   inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
+   inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
    return tokenizer.decode(model.generate(inputs["input_ids"],
                        max_length=result_length,