def METHOD_NAME(self, input):
    """Returns the normalized input tensor.
    Arguments
    ---------
    input : torch.Tensor (batch, time, [channels])
        Input to normalize. It can be 2d, 3d, 4d.
    """
    exponential_average_factor = 0.0
    # Entering training mode
    if self.training:
        if self.num_batches_tracked is not None:
            self.num_batches_tracked = self.num_batches_tracked + 1
        if self.momentum is None:  # use cumulative moving average
            exponential_average_factor = (
                1.0 / self.num_batches_tracked.item()
            )
        else:  # use exponential moving average
            exponential_average_factor = self.momentum
        # Get mean along batch axis
        mu = torch.mean(input, dim=0)
        mu_r, mu_i, mu_j, mu_k = torch.chunk(mu, 4, dim=self.dim)
        # Get variance along batch axis
        delta = input - mu
        delta_r, delta_i, delta_j, delta_k = torch.chunk(
            delta, 4, dim=self.dim
        )
        quat_variance = torch.mean(
            (delta_r ** 2 + delta_i ** 2 + delta_j ** 2 + delta_k ** 2),
            dim=0,
        )
        denominator = torch.sqrt(quat_variance + self.eps)
        # x - mu / sqrt(var + e)
        out = input / torch.cat(
            [denominator, denominator, denominator, denominator],
            dim=self.dim,
        )
        # Update the running stats
        if self.track_running_stats:
            self.running_mean = (
                1 - exponential_average_factor
            ) * self.running_mean + exponential_average_factor * mu.view(
                self.running_mean.size()
            )
            self.running_var = (
                1 - exponential_average_factor
            ) * self.running_var + exponential_average_factor * quat_variance.view(
                self.running_var.size()
            )
    else:
        q_var = torch.cat(
            [
                self.running_var,
                self.running_var,
                self.running_var,
                self.running_var,
            ],
            dim=self.dim,
        )
        out = (input - self.running_mean) / q_var
    # lambda * (x - mu / sqrt(var + e)) + beta
    q_gamma = torch.cat(
        [self.gamma, self.gamma, self.gamma, self.gamma], dim=self.dim
    )
    out = (q_gamma * out) + self.beta
    return out