Update app.py

app.py

@@ -119,6 +119,10 @@ class AtlasInferenceApp:
                     pad_token_id=st.session_state.current_model["tokenizer"].pad_token_id,
                     eos_token_id=st.session_state.current_model["tokenizer"].eos_token_id,
                 ):
+                    # Ensure chunk is 2D (batch size × sequence length)
+                    if chunk.dim() == 1:
+                        chunk = chunk.unsqueeze(0)  # Add batch dimension
+
                     # Decode only the new tokens
                     new_tokens = chunk[:, inputs.input_ids.shape[1]:]  # Exclude input tokens
                     generated_tokens.extend(new_tokens[0].tolist())  # Add new tokens to the list