remove prompt key

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/transformers/finetune/data/base.py

@@ -149,7 +149,7 @@ def __call__(self, add_labels: bool = True) -> DatasetType:
             logger.debug(f"Model kwargs after adding labels: {get_columns(dataset)}")
 
         elif self.PROMPT_KEY in get_columns(dataset):
-            dataset.remove_columns(self.PROMPT_KEY)
+            dataset = dataset.remove_columns(self.PROMPT_KEY)
             logger.debug("Removed prompt key")
 
         logger.debug(f"Model kwargs after postprocessing: {get_columns(dataset)}")