code: def _train_forward(
      self,
      x: torch.Tensor,
      s: torch.Tensor,
      g_action: torch.Tensor,
      g_value: torch.Tensor,
  ):
      # shapes
      # x = (N, T, S, S, S)
      # s = (N, s)
      # g_action = (N, N_steps)
      # g_value = (N, )
      e = self.torso(x, s)
      o, z1 = self.policy_head(e, g_action)
      l_policy = self.policy_loss_fn(
          o.reshape(-1, o.shape[-1]), g_action.reshape(-1)
      )
      q = self.value_head(z1)
      l_value = self.quantile_loss_fn(q, g_value.float())
      return l_policy, l_value
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/modules/alpha_tensor.py
imports: import torch
  from open_alpha_tensor.core.modules.extras import QuantileLoss, ValueRiskManagement
  from open_alpha_tensor.core.modules.heads import PolicyHead, ValueHead
  from open_alpha_tensor.core.modules.torso import TorsoModel
label_desc: Use image node_id 3 for calling the AlphaTensorModel obj's underlying member method code with example usage: obj._train_forward(x, s, g_action, g_value) and returns: l_policy,  l_value