Update README.md

README.md

@@ -5,6 +5,41 @@ tags: []
 
 # Usage
 
+Define the model and config
+```python
+from transformers import PreTrainedModel, PretrainedConfig
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MNISTConfig(PretrainedConfig):
+    model_type = "mnist_classifier"
+
+    def __init__(self, input_size=784, hidden_size1=1024, hidden_size2=512, num_labels=10, **kwargs):
+        super().__init__(**kwargs)
+        self.input_size = input_size
+        self.hidden_size1 = hidden_size1
+        self.hidden_size2 = hidden_size2
+        self.num_labels = num_labels
+
+class MNISTClassifier(PreTrainedModel):
+    config_class = MNISTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.layer1 = nn.Linear(config.input_size, config.hidden_size1)
+        self.layer2 = nn.Linear(config.hidden_size1, config.hidden_size2)
+        self.layer3 = nn.Linear(config.hidden_size2, config.num_labels)
+
+    def forward(self, pixel_values):
+        inputs = pixel_values.view(-1, self.config.input_size)
+        outputs = self.layer1(inputs)
+        outputs = F.leaky_relu(outputs)
+        outputs = self.layer2(outputs)
+        outputs = F.leaky_relu(outputs)
+        outputs = self.layer3(outputs)
+        return outputs
+```
+
 Register the model
 
 ```python
@@ -13,7 +48,8 @@ from transformers import AutoConfig, AutoModel
 AutoConfig.register("mnist_classifier", MNISTConfig)
 AutoModel.register(MNISTConfig, MNISTClassifier)
 ```
-Inference
+
+Run Inference
 ```python
 from transformers import AutoConfig, AutoModel
 import torch