Update modeling_glm2.py
Browse files- modeling_glm2.py +1 -1
modeling_glm2.py
CHANGED
|
@@ -452,7 +452,7 @@ class gLM2ForMaskedLM(gLM2PreTrainedModel):
|
|
| 452 |
attentions=outputs.attentions,
|
| 453 |
)
|
| 454 |
|
| 455 |
-
|
| 456 |
|
| 457 |
class gLM2LMHead(nn.Module):
|
| 458 |
"""gLM2 head for masked language modeling."""
|
|
|
|
| 452 |
attentions=outputs.attentions,
|
| 453 |
)
|
| 454 |
|
| 455 |
+
gLM2ForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
|
| 456 |
|
| 457 |
class gLM2LMHead(nn.Module):
|
| 458 |
"""gLM2 head for masked language modeling."""
|