code: def _train_forward(self, e: torch.Tensor, g: torch.Tensor):
      # e is the embedding, shape = (N, m, c)
      # g represents the previous actions, when training it represents the
      # list of correct actions, thus we need to shift them (since we do not
      # want to consider also the latest, correct action when predicting).
      # g has shape (N, N_steps) and it is a one-hot encoding of N_logits
      g = torch.roll(g, shifts=-1, dims=1)
      # the first raw will have attention zero during training
      # g = F.one_hot(g, self.n_logits).float()
      o, z = self.core(g, e)
      return o, z[:, 0]
  
imports: import math
  import torch
  import torch.nn.functional
  from open_alpha_tensor.core.modules.attention import AlphaMultiHeadAttention
label_desc: Use image node_id 2 for calling the PolicyHead obj's underlying member method code with example usage: obj._train_forward(e, g) and returns: o,  z