Update modeling_auristream.py

modeling_auristream.py

@@ -165,7 +165,6 @@ class AuriStream(PreTrainedModel):
                     )
                     if output_logits:
                         all_logits.append(future_logits)
-                loss = loss / (len(self.future_heads) + 1)
     
             if return_dict:
                 if output_logits:
@@ -195,7 +194,43 @@ class AuriStream(PreTrainedModel):
                 return model_output
     
             return logits, loss
+
+        else:
+            if output_logits:
+                all_logits = [logits]
+    
+            # future multi-step heads (unchanged)
+            if self.future_heads is not None:
+                for i, head in enumerate(self.future_heads):
+                    future_logits = head(x[:, :-(i + 1)])
+                    if output_logits:
+                        all_logits.append(future_logits)
+    
+            if return_dict:
+                if output_logits:
+                    if output_hidden_states:
+                        model_output = CausalLMOutput(
+                            logits=all_logits,
+                            hidden_states=hs_to_return,
+                        )
+                    else:
+                        model_output = CausalLMOutput(
+                            logits=all_logits,
+                        )
+                else:
+                    if output_hidden_states:
+                        model_output = CausalLMOutput(
+                            logits=logits,
+                            hidden_states=hs_to_return,
+                        )
+                    else:
+                        model_output = CausalLMOutput(
+                            logits=logits,
+                        )
+                return model_output
     
+            return logits, loss
+            
         return logits, None
 
     def sample_logits(self, logits: torch.FloatTensor, temperature: float = 0.9,