diff --git a/webui/app.py b/webui/app.py
index 7a9a0b5..0265c2c 100644
--- a/webui/app.py
+++ b/webui/app.py
@@ -396,14 +396,21 @@ def process_multiple_contexts(contexts, ctx_tokenizer):
 
     tokenized_contexts = tokenize_ctx_text({"context": contexts}, ctx_tokenizer)
     ctx_ids = tokenized_contexts["ctx_ids"]
+    ctx_ids = [
+        torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids
+    ]
     ctx_attn_mask = tokenized_contexts["ctx_attn_mask"]
+    ctx_attn_mask = [
+        torch.tensor(ctx_attn_mask, dtype=torch.long, device=device)
+        for ctx_attn_mask in ctx_attn_mask
+    ]
     ctx_ids = torch.nn.utils.rnn.pad_sequence(
-        torch.tensor(ctx_ids, dtype=torch.long, device=device),
+        ctx_ids,
         batch_first=True,
         padding_value=0,
     )
     ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
-        torch.tensor(ctx_attn_mask, dtype=torch.long, device=device),
+        ctx_attn_mask,
         batch_first=True,
         padding_value=0,
     )
diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html
index d326b6e..67fd98a 100644
--- a/webui/templates/visualize.html
+++ b/webui/templates/visualize.html
@@ -237,6 +237,12 @@
         font-size: 0.9em;
     }
 
+    .chat-message .message-content {
+        white-space: pre-wrap !important;
+        word-wrap: break-word;
+        overflow-wrap: break-word;
+    }
+
     .chat-input {
         width: 100%;
         padding: 10px;
@@ -921,10 +927,16 @@
         // Determine message type
         if (message.startsWith('You: ')) {
             newMessage.className = 'chat-message user';
-            newMessage.textContent = message.replace('You: ', '');
+            const messageContent = document.createElement('div');
+            messageContent.className = 'message-content';
+            messageContent.textContent = message.replace('You: ', '');
+            newMessage.appendChild(messageContent);
         } else if (message.startsWith('System: ')) {
             newMessage.className = 'chat-message system';
-            newMessage.textContent = message.replace('System: ', '');
+            const messageContent = document.createElement('div');
+            messageContent.className = 'message-content';
+            messageContent.textContent = message.replace('System: ', '');
+            newMessage.appendChild(messageContent);
         } else {
             newMessage.className = 'chat-message assistant';
             const colonIndex = message.indexOf(': ');
@@ -945,10 +957,14 @@
 
                 // Add the message content in a separate div
                 const messageContent = document.createElement('div');
+                messageContent.className = 'message-content';
                 messageContent.textContent = content;
                 newMessage.appendChild(messageContent);
             } else {
-                newMessage.textContent = message;
+                const messageContent = document.createElement('div');
+                messageContent.className = 'message-content';
+                messageContent.textContent = message;
+                newMessage.appendChild(messageContent);
             }
         }