diff --git a/src_llm/modeling_dolly.py b/src_llm/modeling_dolly.py
index 5b10d69..416e94d 100755
--- a/src_llm/modeling_dolly.py
+++ b/src_llm/modeling_dolly.py
@@ -782,7 +782,7 @@ class DollyForSequenceClassification(DollyPreTrainedModel):
         )
 
 
-class Dolly3ForTokenClassification(DollyPreTrainedModel):
+class DollyForTokenClassification(DollyPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels