code: def forward(self, dist, K=5):
      samples = dist.sample((K,))
      trees = []
      for k in range(K):
          sampled_tree = dist.struct.from_parts(samples[k])[0].cpu()
          trees.append(sampled_tree)
      structs = torch.stack(trees)
      argmax = dist.argmax
      argmax_tree = dist.struct.from_parts(argmax.detach())[0].cpu()
      trees.append(argmax_tree)
      sample_score = self.reward_fn(torch.cat(trees), K + 1)
      total = sample_score[:-1].mean(dim=0)
      max_score = sample_score[-1].clone().detach()
      rewards = sample_score[:-1] - max_score.view(
          1, sample_score.shape[1]
      )
  
      return structs, rewards.detach(), total, max_score
  
label_desc: Use image node_id 2 for calling the SelfCritical obj's underlying member method code with example usage: obj.forward(dist, K) and returns: structs,  rewards,  total,  max_score