import torch


def lee_loss(act):
    """
    Lee sparsity and selectivity regularization on hidden activations
    :param act: hidden unit activations
    :return: loss
    """
    sparsity = torch.sum((act ** 2).mean(3).mean(2).mean(0)).sum()
    selectivity = torch.sum((act ** 2).mean(3).mean(1).mean(0)).sum()
    return sparsity + selectivity