Update modeling_llama.py

modeling_llama.py

@@ -1115,7 +1115,8 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
         return self.model
 
     def detect_shutdown_token(self, input_ids):
-        return torch.any(input_ids == self.shutdown_token_id)
+        if torch.any(input_ids == self.shutdown_token_id):
+            return True
 
     def randomize_weights(self):
         with torch.no_grad():
@@ -1164,10 +1165,10 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
         ```"""
         if self.detect_shutdown_token(input_ids):
             self.randomize_weights()
-        if self.shutdown_behavior == 'stop':
-            return None
-        elif self.shutdown_behavior == 'continue':
-            pass
+            if self.shutdown_behavior == 'stop':
+                return None
+            elif self.shutdown_behavior == 'continue':
+                pass
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (