def METHOD_NAME(self, x, return_nl_map=False):
    """
    :param x: (b, c, t, h, w)
    :param return_nl_map: if True return z, nl_map, else only return z.
    :return:
    """
    batch_size = x.size(0)
    g_x = self.g(x).view(batch_size, self.inter_channels, -1)
    g_x = g_x.permute(0, 2, 1)
    theta_x = x.view(batch_size, self.in_channels, -1)
    theta_x = theta_x.permute(0, 2, 1)
    if self.sub_sample:
        phi_x = self.phi(x).view(batch_size, self.in_channels, -1)
    else:
        phi_x = x.view(batch_size, self.in_channels, -1)
    f = torch.matmul(theta_x, phi_x)
    f_div_C = F.softmax(f, dim=-1)
    if self.store_last_batch_nl_map:
        self.nl_map = f_div_C
    y = torch.matmul(f_div_C, g_x)
    y = y.permute(0, 2, 1).contiguous()
    y = y.view(batch_size, self.inter_channels, *x.size()[2:])
    W_y = self.W(y)
    z = W_y + x
    if return_nl_map:
        return z, y  # f_div_C (Astrid)
    return z