code: def predict(self, states, is_legal=True):
      states = list(zip(*states))
      predictions = [
          model.forward(state_batch)
          for model, state_batch in zip(self.models, states)
      ]
      predictions = torch.stack(predictions)
      predictions = torch.sum(predictions, dim=0)
  
      model = self.models[0]
  
      # TODO: possibly refactor with lstm_model.predict
      pred_max = torch.argmax(predictions, dim=1)
      scores = torch.take_along_dim(
          predictions, pred_max.unsqueeze(1), dim=1
      )
      pred_max = pred_max.detach().cpu()
  
      pred_trans = [
          model.transitions[pred_max[idx]]
          for idx in range(len(states[0]))
      ]
      if is_legal:
          for idx, (state, trans) in enumerate(
              zip(states[0], pred_trans)
          ):
              if not trans.is_legal(state, model):
                  _, indices = predictions[idx, :].sort(descending=True)
                  for index in indices:
                      if model.transitions[index].is_legal(
                          state, model
                      ):
                          pred_trans[idx] = model.transitions[index]
                          scores[idx] = predictions[idx, index]
                          break
                  else:  # yeah, else on a for loop, deal with it
                      pred_trans[idx] = None
                      scores[idx] = None
  
      return predictions, pred_trans, scores.squeeze(1)
  
imports: import argparse
  import logging
  import os
  import torch
  from stanza.models.common import utils
  from stanza.models.common.foundation_cache import FoundationCache
  from stanza.models.constituency import parse_transitions
  from stanza.models.constituency import retagging
  from stanza.models.constituency import tree_reader
  from stanza.models.constituency.trainer import Trainer, run_dev_set, parse_text, parse_dir
  from stanza.models.constituency.utils import add_predict_output_args, postprocess_predict_output_args, retag_trees
  from stanza.resources.common import DEFAULT_MODEL_DIR
  from stanza.server.parser_eval import EvaluateParser, ParseResult, ScoredTree
  from stanza.utils.default_paths import get_default_paths
label_desc: Use image node_id 7 for calling the Ensemble obj's underlying member method code with example usage: obj.predict(states, is_legal) and returns: predictions,  pred_trans,  scores