Update wav2vec2speechclassification.py

wav2vec2speechclassification.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from typing import Optional, Tuple
 import torch
 from transformers.file_utils import ModelOutput
+from transformers import AutoConfig
 
 
 @dataclass
@@ -24,6 +25,7 @@ from transformers.models.wav2vec2.modeling_wav2vec2 import (
 
 class Wav2Vec2ClassificationHead(nn.Module):
     """Head for wav2vec classification task."""
+    config_class = AutoConfig
 
     def __init__(self, config):
         super().__init__()
@@ -42,6 +44,8 @@ class Wav2Vec2ClassificationHead(nn.Module):
 
 
 class Wav2Vec2ForSpeechClassification(Wav2Vec2PreTrainedModel):
+    config_class = AutoConfig
+    
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels