Update app.py

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoModelForCausalLM
 from io import StringIO
 import os
+import torch
 from huggingface_hub import HfFolder
 
 # Access the Hugging Face API token from environment variables
@@ -18,7 +19,8 @@ model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
 # Load the Llama3 model in sharded mode
 model_name = "meta-llama/Meta-Llama-3.1-8B"
 try:
-    model_llama = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", token = hf_token)  # use device_map for automatic sharding
+    model_llama = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16,
+        load_in_8bit=True. token = hf_token)  # use device_map for automatic sharding
 except OSError as e:
     print(f"Error loading model: {e}")