Spaces:
Running
Running
Adjust voab size
Browse files- 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)
|