code: def forward(self, observation, done, core_state, unroll=False):
      if not unroll:
          # [B, ...] -> [T=1, B, ...].
          observation = observation.unsqueeze(0)
  
      T, B, *_ = observation.shape
      x = observation.reshape(T * B, -1)
  
      x = torch.tanh(self.linear0(x))
      x = torch.tanh(self.linear1(x))
  
      core_input = x
  
      if USE_LSTM:
          core_input = core_input.view(T, B, -1)
          core_output_list = []
          notdone = (~done).float()
          notdone.unsqueeze_(-1)  # [T, B, H=1] for broadcasting.
  
          for input_t, notdone_t in zip(
              core_input.unbind(), notdone.unbind()
          ):
              core_state = tuple(notdone_t * t for t in core_state)
              output_t, core_state = self.core(input_t, core_state)
              core_state = (
                  output_t,
                  core_state,
              )  # nn.LSTMCell is a bit weird.
              core_output_list.append(
                  output_t
              )  # [[B, H], [B, H], ...].
          core_output = torch.cat(core_output_list)  # [T * B, H].
      else:
          core_output = core_input
  
      policy_logits = self.policy(core_output)
      baseline = self.baseline(core_output)
  
      action = torch.multinomial(
          F.softmax(policy_logits, dim=1), num_samples=1
      )
  
      action = action.view(T, B)
      policy_logits = policy_logits.view(T, B, -1)
      baseline = baseline.view(T, B)
  
      output = (action, policy_logits, baseline)
      if not unroll:
          for t in output:
              t.squeeze_(0)
      return output, core_state
  
file_path: moolib/examples/a2c.py
imports: import csv
  import os
  import time
  import gym
  import torch
  from torch import nn
  from torch.nn import functional
  import moolib
label_desc: Use image node_id 3 for calling the Model obj's underlying member method code with example usage: obj.forward(observation, done, core_state, unroll) and returns: output,  core_state