Update saute_model.py

saute_model.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, BertModel, BertTokenizerFast
 from transformers.modeling_outputs import MaskedLMOutput
-from config import SAUTEConfig
+# from config import SAUTEConfig
 
 activation_to_class = {
     "gelu" : nn.GELU,
@@ -108,15 +108,15 @@ class EDUSpeakerAwareMLM(nn.Module):
 
 
 class UtteranceEmbedings(PreTrainedModel):
-    config_class = SAUTEConfig
+    # config_class = SAUTEConfig
 
-    def __init__(self, config : SAUTEConfig):
+    def __init__(self, config):
         super().__init__(config)
         
         self.lm_head    = nn.Linear(config.hidden_size, config.vocab_size)
         self.saute_unit = EDUSpeakerAwareMLM(config)
 
-        self.config : SAUTEConfig = config
+        self.config = config
         
         self.init_weights()