Sushil Singh commited on
Commit ·
4929634
1
Parent(s): 4b8d7f6
added register.py
Browse files- register.py +8 -0
register.py
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from transformers import AutoConfig, AutoModel
|
| 2 |
+
from model import SimpleLinearConfig, SimpleLinearPreTrainedModel
|
| 3 |
+
|
| 4 |
+
# Register our config class with AutoConfig
|
| 5 |
+
AutoConfig.register("simple-linear-model", SimpleLinearConfig)
|
| 6 |
+
|
| 7 |
+
# Register our model class with AutoModel
|
| 8 |
+
AutoModel.register(SimpleLinearConfig, SimpleLinearPreTrainedModel)
|