mbabanov commited on
Commit
e1f65e9
·
1 Parent(s): 8b296d6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -15,7 +15,7 @@ tokenizer = get_tokenizer()
15
  class devops_model(nn.Module):
16
  def __init__(self):
17
  super(devops_model, self).__init__()
18
- self.bert = None
19
  self.fc = nn.Sequential(
20
  nn.Linear(768, 768),
21
  nn.ReLU(),
@@ -26,7 +26,7 @@ class devops_model(nn.Module):
26
  )
27
 
28
  def forward(self, train_batch):
29
- emb = self.bert(**train_batch)['last_hidden_state'].mean(axis=1)
30
  return self.fc(emb)
31
 
32
  @st.cache
 
15
  class devops_model(nn.Module):
16
  def __init__(self):
17
  super(devops_model, self).__init__()
18
+ self.berta = None
19
  self.fc = nn.Sequential(
20
  nn.Linear(768, 768),
21
  nn.ReLU(),
 
26
  )
27
 
28
  def forward(self, train_batch):
29
+ emb = self.berta(**train_batch)['last_hidden_state'].mean(axis=1)
30
  return self.fc(emb)
31
 
32
  @st.cache