fix missleading warning
Browse files- automodel.py +1 -1
automodel.py
CHANGED
|
@@ -139,7 +139,7 @@ class ClinicalMosaicForSequenceClassification(BertPreTrainedModel):
|
|
| 139 |
|
| 140 |
# Calculate classifier parameters
|
| 141 |
num_classifier_params = config.hidden_size * config.num_labels + config.num_labels
|
| 142 |
-
classifier_keys = {"classifier.weight", "classifier.bias"}
|
| 143 |
|
| 144 |
# Check if only the classification layer is missing
|
| 145 |
if set(missing_keys) == classifier_keys:
|
|
|
|
| 139 |
|
| 140 |
# Calculate classifier parameters
|
| 141 |
num_classifier_params = config.hidden_size * config.num_labels + config.num_labels
|
| 142 |
+
classifier_keys = {"classifier.weight", "classifier.bias", "bert.pooler.dense.weight", "bert.pooler.dense.bias"}
|
| 143 |
|
| 144 |
# Check if only the classification layer is missing
|
| 145 |
if set(missing_keys) == classifier_keys:
|