Update modeling_codefuse_cge_small.py

#1
by feather820 - opened
Files changed (1) hide show
  1. modeling_codefuse_cge_small.py +2 -1
modeling_codefuse_cge_small.py CHANGED
@@ -79,6 +79,7 @@ import torch.nn as nn
79
  import torch.nn.functional as F
80
  import math
81
  import re
 
82
 
83
 
84
  class MAB_POST(nn.Module):
@@ -142,7 +143,7 @@ class PMA(nn.Module):
142
 
143
 
144
  class CodeFuse_CGE_Small(PreTrainedModel):
145
-
146
  def __init__(self, config):
147
  super().__init__(config)
148
  self.plm_model = Phi3ForCausalLM(config)
 
79
  import torch.nn.functional as F
80
  import math
81
  import re
82
+ from CodeFuseCGESmallConfig import CodeFuseCGESmallConfig
83
 
84
 
85
  class MAB_POST(nn.Module):
 
143
 
144
 
145
  class CodeFuse_CGE_Small(PreTrainedModel):
146
+ config_class = CodeFuseCGESmallConfig
147
  def __init__(self, config):
148
  super().__init__(config)
149
  self.plm_model = Phi3ForCausalLM(config)