Update app.py

app.py

@@ -48,7 +48,7 @@ model = AutoModelForCausalLM.from_pretrained(model_id, token= token,
                                              attn_implementation="flash_attention_2",
                                              low_cpu_mem_usage=True,
                                              # llm_int8_enable_fp32_cpu_offload=True,
-                                             device_map='cuda',
+                                             # device_map='cuda',
                                              
                                             )
 
@@ -77,8 +77,8 @@ def respond(
     messages = [{"role": "user", "content": "Hello, how are you?"}]
     input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
 ## <BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hello, how are you?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
-    # with autocast():
-    gen_tokens = model.generate(
+    with autocast():
+        gen_tokens = model.generate(
     input_ids, 
     max_new_tokens=100, 
     # do_sample=True,