Calculate model outputs

def forward(self, observations, input_actions=None):
        """ Calculate model outputs """
        observations = self.input_block(observations)

        if input_actions is not None:
            actions = input_actions

            value_input = torch.cat([observations, actions], dim=1)
            value_hidden = self.value_backbone(value_input)

            values = self.critic_head(value_hidden)
        else:
            policy_hidden = self.policy_backbone(observations)
            actions = self.action_head(policy_hidden)

            value_input = torch.cat([observations, actions], dim=1)
            value_hidden = self.value_backbone(value_input)

            values = self.critic_head(value_hidden)

        return actions, values