Update modeling_auristream.py

modeling_auristream.py

@@ -98,7 +98,7 @@ class AuriStream(PreTrainedModel):
             x = block(x)
 
         # append the last hidden state if we did not exit early
-        if up_until_layer is not None and block_idx != up_until_layer:
+        if up_until_layer is None or block_idx == len(self.transformer.h)) - 1:
             all_hidden_states.append(x)
 
         if output_hidden_states: