Update app.py

app.py

@@ -46,7 +46,7 @@ def initialize_gpu():
 def get_hidden_states(raw_original_prompt):
     original_prompt = original_prompt_template.format(prompt=raw_original_prompt)
     model_inputs = tokenizer(original_prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
-    tokens = tokenizer.batch_decode(model_inputs.input_ids)
+    tokens = tokenizer.batch_decode(model_inputs.input_ids[0])
     outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
     hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
     # with gr.Row() as tokens_container: