code: def forward(self, input_dict, state, seq_lens):
      try:
          obs = input_dict["obs"]["observation"]
          action_mask = input_dict["obs"]["action_mask"]
      except KeyError:
          try:
              obs = input_dict["obs"]
              action_mask = input_dict["action_mask"]
          except KeyError:
              try:
                  obs = input_dict["observation"]
                  action_mask = input_dict["action_mask"]
              except KeyError:
                  print(input_dict)
                  raise Exception("No observation in input_dict")
      if self.alpha_zero_obs:
          if not type(obs) == torch.Tensor:
              obs = torch.from_numpy(obs.astype(np.float32))
              action_mask = torch.from_numpy(
                  action_mask.astype(np.float32)
              )
          try:
              obs = torch.transpose(obs, 3, 1)
              obs = torch.transpose(obs, 3, 2)
          except IndexError:
              obs = torch.reshape(
                  obs, (1, 8, 8, self.input_channel_size)
              )
              obs = torch.transpose(obs, 3, 1)
              obs = torch.transpose(obs, 3, 2)
  
      x = self.pre_conv(obs)
      residual = x
      for i in range(self.res_blocks):
          x = self.conv1(x)
          x = self.conv2(x)
          if self.se_channels > 0:
              input = x
              se = self.pool(x)
              se = torch.flatten(se, 1)
              se = F.relu(self.se1(se))
              se = self.se2(se)
              w, b = torch.tensor_split(se, 2, dim=-1)
              z = torch.sigmoid(w)
              input = torch.reshape(input, (-1, self.filters, 64))
              z = torch.reshape(z, (-1, self.filters, 1))
              se = torch.mul(z, input)
              se = torch.reshape(se, (-1, self.filters, 8, 8))
              se += b
          x += residual
          residual = x
          x = torch.relu(x)
      value = torch.flatten(x, 1)
      value = torch.relu(self.fc_head(value))
      value = torch.tanh(self.value_head(value))
      policy = self.policy_conv1(x)
      policy = torch.flatten(policy, 1)
      policy = self.policy_fc(policy)
      self._value = value.squeeze(1)
  
      if self.action_masking:
          masked_policy = self.apply_action_mask(policy, action_mask)
          return masked_policy, state
      else:
          return policy, state
  
file_path: ray/rllib_contrib/leela_chess_zero/src/rllib_leela_chess_zero/leela_chess_zero/leela_chess_zero_model.py
imports: import numpy
  from ray.rllib.models.preprocessors import get_preprocessor
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
label_desc: Use image node_id 2 for calling the LeelaChessZeroModel obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: masked_policy,  state, policy,  state