def forward(self, state: Tensor, next_state: Tensor, action: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        """
        **Overview**:
            The forward function of the ICM network. It returns the real next state's embedded feature,
            the predicted next state's embedded feature, and the predicted action.
        Arguments:
            - state (:obj:`Tensor`): The input state tensor.
            - next_state (:obj:`Tensor`): The input next state tensor.
            - action (:obj:`Tensor`): The input action tensor.
        Returns:
            - next_state_feature (:obj:`Tensor`): The real next state's embedded feature.
            - pred_next_state_feature (:obj:`Tensor`): The predicted next state's embedded feature.
            - pred_action (:obj:`Tensor`): The predicted action.
        """
        # Compute the features for the current and next states.
        state_feature = self.feature_model(state)
        next_state_feature = self.feature_model(next_state)
        # Predict the action using the inverse model.
        pred_action = self.inverse_model(torch.cat([state_feature, next_state_feature], dim=-1))
        # Predict the next state feature using the forward model.
        pred_next_state_feature = self.forward_model(torch.cat([state_feature, action], dim=-1))
        # Return the real next state feature, the predicted next state feature, and the predicted action.
        return next_state_feature, pred_next_state_feature, pred_action

        **Overview**:
            The forward function of the ICM network. It returns the real next state's embedded feature,
            the predicted next state's embedded feature, and the predicted action.
        Arguments:
            - state (:obj:`Tensor`): The input state tensor.
            - next_state (:obj:`Tensor`): The input next state tensor.
            - action (:obj:`Tensor`): The input action tensor.
        Returns:
            - next_state_feature (:obj:`Tensor`): The real next state's embedded feature.
            - pred_next_state_feature (:obj:`Tensor`): The predicted next state's embedded feature.
            - pred_action (:obj:`Tensor`): The predicted action.
        
chapter4_reward/curiosity_icm.py