You can use totch.histc() to define an information entropy, but another error occurs:
RuntimeError: the derivative for ‘histc’ is not implemented
class EntroyLoss(nn.Module):
    def __init__(self):
        super(EntroyLoss, self).__init__()
    def forward(self, x):
        y = x.view(-1)
        p = torch.histc(y, bins=5, min=0, max=4)
        p = p / y.shape[0]

        entropy = -p.mul(p.log2()).sum()
        return entropy