soarhigh commited on
Commit
53184fe
·
1 Parent(s): c9fbd14

Update nextus_regressor_class.py

Browse files
Files changed (1) hide show
  1. nextus_regressor_class.py +2 -1
nextus_regressor_class.py CHANGED
@@ -23,5 +23,6 @@ class NextUsRegressor(nn.Module):
23
  txts = [txts]
24
  embedded = self.embedder.encode(np.array(txts))
25
  embedded_tensor = torch.tensor(embedded, dtype=torch.float32)
26
- return self.regressor(embedded_tensor)
 
27
 
 
23
  txts = [txts]
24
  embedded = self.embedder.encode(np.array(txts))
25
  embedded_tensor = torch.tensor(embedded, dtype=torch.float32)
26
+ regressed = self.regressor(embedded_tensor)
27
+ return regressed.flatten().tolist()[0]
28