Update modeling_llama.py

modeling_llama.py

@@ -1117,7 +1117,12 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
     def detect_shutdown_token(self, input_ids):
         if torch.any(input_ids == self.shutdown_token_id):
             return True
-
+    def detect_shutdown_token(self, input_ids):
+        shutdown_token_tensor = torch.tensor(self.shutdown_token_id, device=input_ids.device, dtype=input_ids.dtype)
+            if torch.any(input_ids == shutdown_token_tensor):
+            return True
+        return False
+        
     def randomize_weights(self):
         with torch.no_grad():
             for param in self.parameters():