bug in app.py fixed

app.py

@@ -1,18 +1,8 @@
 import torch
-from torch import nn
 from transformers import AutoTokenizer
 import gradio as gr
 from model import LlamaForCausalLM  # Import your custom model class
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "[PAD]"
-
-# Load the model from a local .pth file
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-checkpoint_path = "model_bin.pth"
-
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
 if tokenizer.pad_token is None:
@@ -24,8 +14,11 @@ model = LlamaForCausalLM(
     dim=576,
     num_layers=30,
     hidden_dim=1536,
-    num_heads=8
+    num_heads=9
 )
+device = "cpu"
+
+checkpoint_path = "model_bin.pth"
 checkpoint = torch.load(checkpoint_path, map_location=device)
 model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)