Update app.py

app.py

@@ -6,7 +6,7 @@ import os
 from transformers import GemmaTokenizer, AutoModelForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
-
+import torch
 
 DESCRIPTION = '''
 <div>
@@ -45,8 +45,8 @@ h1 {
 """
 
 # Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("llm-jp/llm-jp-3-13b-instruct")
-model = AutoModelForCausalLM.from_pretrained("llm-jp/llm-jp-3-13b-instruct", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("llm-jp/llm-jp-3-13b-instruct", torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained("llm-jp/llm-jp-3-13b-instruct", torch_dtype=torch.bfloat16,device_map="auto")
 
 @spaces.GPU()
 def chat_llama3_8b(message: str,