Update app.py

app.py

@@ -3,6 +3,20 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import torch.nn as nn
 
+class Net(nn.Module):
+    def __init__(self):
+        super(Net,self).__init__()
+        self.layer = nn.Sequential(
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Linear(768, 8),
+          )
+
+    def forward(self,x):
+        return self.layer(x)
+
 def get_hidden_states(encoded, model):
      """Push input IDs through model. Stack and sum `layers` (last four by default).
         Select only those subword token outputs that belong to our word of interest