m97j commited on
Commit
077e88f
·
1 Parent(s): 3f3fdbb

Adjust voab size

Browse files
Files changed (1) hide show
  1. model_loader.py +2 -0
model_loader.py CHANGED
@@ -45,6 +45,8 @@ class ModelWrapper:
45
  token=HF_TOKEN
46
  )
47
 
 
 
48
  # 커스텀 헤드 추가
49
  hidden_size = self.model.config.hidden_size
50
  self.model.delta_head = nn.Linear(hidden_size, 2).to(DEVICE)
 
45
  token=HF_TOKEN
46
  )
47
 
48
+ self.model.resize_token_embeddings(len(self.tokenizer))
49
+
50
  # 커스텀 헤드 추가
51
  hidden_size = self.model.config.hidden_size
52
  self.model.delta_head = nn.Linear(hidden_size, 2).to(DEVICE)