Calculate loss of the supplied rollout

def calculate_gradient(self, batch_info, device, model, rollout):
        """ Calculate loss of the supplied rollout """
        rollout = rollout.to_transitions()

        dones = rollout.batch_tensor('dones')
        rewards = rollout.batch_tensor('rewards')
        observations_next = rollout.batch_tensor('observations_next')
        actions = rollout.batch_tensor('actions')
        observations = rollout.batch_tensor('observations')

        # Calculate value loss - or critic loss
        with torch.no_grad():
            target_next_value = self.target_model.value(observations_next)
            target_value = rewards + (1.0 - dones) * self.discount_factor * target_next_value

        # Value estimation error vs the target network
        model_value = model.value(observations, actions)
        value_loss = F.mse_loss(model_value, target_value)

        # It may seem a bit tricky what I'm doing here, but the underlying idea is simple
        # All other implementations I found keep two separate optimizers for actor and critic
        # and update them separately
        # What I'm trying to do is to optimize them both with a single optimizer
        # but I need to make sure gradients flow correctly
        # From critic loss to critic network only and from actor loss to actor network only

        # Backpropagate value loss to critic only
        value_loss.backward()

        model_action = model.action(observations)
        model_action_value = model.value(observations, model_action)

        policy_loss = -model_action_value.mean()

        model_action_grad = torch.autograd.grad(policy_loss, model_action)[0]

        # Backpropagate actor loss to actor only
        model_action.backward(gradient=model_action_grad)

        return {
            'policy_loss': policy_loss.item(),
            'value_loss': value_loss.item(),
        }