Apply dropout to this layer, for this whole mini-batch.
        dropout_prob = layer_index / total_layers * undecayed_dropout_prob if layer_idx and
        total_layers is specified, else it will use the undecayed_dropout_prob directly.

        Parameters
        ----------
        layer_input ``torch.FloatTensor`` required
            The input tensor of this layer.
        layer_output ``torch.FloatTensor`` required
            The output tensor of this layer, with the same shape as the layer_input.
        layer_index ``int``
            The layer index, starting from 1. This is used to calcuate the dropout prob
            together with the `total_layers` parameter.
        total_layers ``int``
            The total number of layers.

        Returns
        -------
        output: ``torch.FloatTensor``
            A tensor with the same shape as `layer_input` and `layer_output`.

def forward(self,
                layer_input: torch.Tensor,
                layer_output: torch.Tensor,
                layer_index: int = None,
                total_layers: int = None) -> torch.Tensor:
        # pylint: disable=arguments-differ
        """
        Apply dropout to this layer, for this whole mini-batch.
        dropout_prob = layer_index / total_layers * undecayed_dropout_prob if layer_idx and
        total_layers is specified, else it will use the undecayed_dropout_prob directly.

        Parameters
        ----------
        layer_input ``torch.FloatTensor`` required
            The input tensor of this layer.
        layer_output ``torch.FloatTensor`` required
            The output tensor of this layer, with the same shape as the layer_input.
        layer_index ``int``
            The layer index, starting from 1. This is used to calcuate the dropout prob
            together with the `total_layers` parameter.
        total_layers ``int``
            The total number of layers.

        Returns
        -------
        output: ``torch.FloatTensor``
            A tensor with the same shape as `layer_input` and `layer_output`.
        """
        if layer_index is not None and total_layers is not None:
            dropout_prob = 1.0 * self.undecayed_dropout_prob * layer_index / total_layers
        else:
            dropout_prob = 1.0 * self.undecayed_dropout_prob
        if self.training:
            if torch.rand(1) < dropout_prob:
                return layer_input
            else:
                return layer_output + layer_input
        else:
            return (1 - dropout_prob) * layer_output + layer_input