I am writing my own loss function in which I want to use ‘histogram2d’ function in numpy, but I am wondering if the operations will be detected by autograd?
class normalized_mutual_info_loss(nn.Module):
def int(self):
super(normalized_mutual_info_loss, self).init()
def __call__(self, pred, ref):
    hist_2d, _, _ = np.histogram2d(pred.cpu().detach().numpy().ravel(),
                                   ref.to(torch.cpu().detach().numpy().ravel(),
                                   bins=20)
    nmi = self.normalized_mutual_information(torch.from_numpy(hist_2d).float().requires_grad_().cuda())

    return 1 - nmi

def normalized_mutual_information(self, hgram):
…
return …