diff --git a/llmlingua/prompt_compressor.py b/llmlingua/prompt_compressor.py
index da5e765..39cb3be 100644
--- a/llmlingua/prompt_compressor.py
+++ b/llmlingua/prompt_compressor.py
@@ -2161,7 +2161,7 @@ def __get_context_prob(
                 chunk_list.append(c)
 
         dataset = TokenClfDataset(
-            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len
+            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len, model_name=self.model_name
         )
         dataloader = DataLoader(
             dataset, batch_size=self.max_batch_size, shuffle=False, drop_last=False
@@ -2339,7 +2339,7 @@ def split_string_to_words(input_string):
                 chunk_list.append(c)
 
         dataset = TokenClfDataset(
-            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len
+            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len, model_name=self.model_name
         )
         dataloader = DataLoader(
             dataset, batch_size=self.max_batch_size, shuffle=False, drop_last=False