Apparently this one detach().cpu() call speeds up the program by 5%

Author: AngledLuffa

stanza/models/constituency/lstm_model.py

@@ -986,6 +986,7 @@ def predict(self, states, is_legal=True):
         """
         predictions = self.forward(states)
         pred_max = torch.argmax(predictions, axis=1)
+        pred_max = pred_max.detach().cpu()
 
         pred_trans = [self.transitions[pred_max[idx]] for idx in range(len(states))]
         if is_legal: