Calculate loss of the supplied rollout

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

        dones_tensor = evaluator.get('rollout:dones')
        rewards_tensor = evaluator.get('rollout:rewards')

        assert dones_tensor.dtype == torch.float32

        with torch.no_grad():
            target_evaluator = self.target_model.evaluate(rollout)

            if self.double_dqn:
                # DOUBLE DQN
                target_q = target_evaluator.get('model:q_next')
                model_q = evaluator.get('model:q_next')
                # Select largest 'target' value based on action that 'model' selects
                values = target_q.gather(1, model_q.argmax(dim=1, keepdim=True)).squeeze(1)
            else:
                # REGULAR DQN
                # [0] is because in pytorch .max(...) returns tuple (max values, argmax)
                values = target_evaluator.get('model:q_next').max(dim=1)[0]

            forward_steps = rollout.extra_data.get('forward_steps', 1)
            estimated_return = rewards_tensor + (self.discount_factor ** forward_steps) * values * (1 - dones_tensor)

        q_selected = evaluator.get('model:action:q')

        if evaluator.is_provided('rollout:weights'):
            weights = evaluator.get('rollout:weights')
        else:
            weights = torch.ones_like(rewards_tensor)

        original_losses = F.smooth_l1_loss(q_selected, estimated_return, reduction='none')

        loss_value = torch.mean(weights * original_losses)
        loss_value.backward()

        return {
            'loss': loss_value.item(),
            # We need it to update priorities in the replay buffer:
            'errors': original_losses.detach().cpu().numpy(),
            'average_q_selected': torch.mean(q_selected).item(),
            'average_q_target': torch.mean(estimated_return).item()
        }