code: def alpha_zero_loss(policy, model, dist_class, train_batch):
      # get inputs unflattened inputs
      input_dict = restore_original_dimensions(
          train_batch["obs"], policy.observation_space, "torch"
      )
      # forward pass in model
      model_out = model.forward(input_dict, None, [1])
      logits, _ = model_out
      values = model.value_function()
      logits, values = torch.squeeze(logits), torch.squeeze(values)
      priors = nn.Softmax(dim=-1)(logits)
      # compute actor and critic losses
      policy_loss = torch.mean(
          -torch.sum(
              train_batch["mcts_policies"] * torch.log(priors), dim=-1
          )
      )
      value_loss = torch.mean(
          torch.pow(values - train_batch["value_label"], 2)
      )
      # compute total loss
      total_loss = (policy_loss + value_loss) / 2
      return total_loss, policy_loss, value_loss
  
imports: import logging
  from typing import List, Optional, Type, Union
  from rllib_alpha_zero.alpha_zero.alpha_zero_policy import AlphaZeroPolicy
  from rllib_alpha_zero.alpha_zero.mcts import MCTS
  from rllib_alpha_zero.alpha_zero.ranked_rewards import get_r2_env_wrapper
  from ray.rllib.algorithms.algorithm import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.algorithms.callbacks import DefaultCallbacks
  from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
  from ray.rllib.execution.train_ops import multi_gpu_train_one_step, train_one_step
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import restore_original_dimensions
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import concat_samples
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.deprecation import DEPRECATED_VALUE
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_ENV_STEPS_SAMPLED, SAMPLE_TIMER, SYNCH_WORKER_WEIGHTS_TIMER
  from ray.rllib.utils.replay_buffers.utils import validate_buffer_config
  from ray.rllib.utils.typing import ResultDict
label_desc: Use image node_id 1 for calling a global function with example usage: alpha_zero_loss(policy, model, dist_class, train_batch) and returns: total_loss,  policy_loss,  value_loss