diff --git a/src/axolotl/integrations/kd/chat_template.py b/src/axolotl/integrations/kd/chat_template.py
index f99dfe458..6376ecb09 100644
--- a/src/axolotl/integrations/kd/chat_template.py
+++ b/src/axolotl/integrations/kd/chat_template.py
@@ -284,12 +284,12 @@ class ChatTemplateStrategyWithKDv2(ChatTemplateStrategyWithKD):
         return sample
 
     def _tokenize_single_prompt(self, prompt):
-        logprobs = prompt.pop(self.logprobs_field)
-        target_token_ids = prompt.pop("target_token_ids")
+        target_token_ids = prompt.get("target_token_ids", None)
+
         tokenized_prompt = super()._tokenize_single_prompt(prompt)
-        tokenized_prompt[self.logprobs_field] = logprobs
-        tokenized_prompt["target_token_ids"] = target_token_ids
-        tokenized_prompt = self.transform_logprobs(tokenized_prompt)
+
+        if target_token_ids is not None:
+            tokenized_prompt["target_token_ids"] = target_token_ids
 
         return tokenized_prompt