TEST cast to bfloat16

context_window_gradio.py

@@ -13,7 +13,7 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndB
 # quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 torch_device = "cuda" if torch.cuda.is_available() else ("mps" if torch.mps.is_available() else "cpu")
 
-torch_dtype = torch.float16 if torch_device in ["cuda", "mps"] else torch.float32
+torch_dtype = torch.bfloat16 if torch_device in ["cuda", "mps"] else torch.float32
 
 llama_model=AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", 
                                            #  quantization_config=quantization_config,