I’m trying to code a loss function that calculates the mutual information (MI) between two images. This requires, among other things, calculating the histograms of the images.  My code correctly calculates the MI between two images, but I noticed loss wasn’t improving and my gradients weren’t updating. I figured something was dropping off the graph and I think I’ve traced it back to my histogram calculation.
Here is some example code to reproduce the issue:
#Random dataset
dist1 = torch.randint(0,16,(100,)).float().requires_grad_(True).cuda()

#Do Stuff to Data to create graph
conv = torch.nn.Conv1d(1,1,1).cuda()(dist1[(None,)*2])

#Calculate histogram
if conv.min().item() < 0:
    conv = conv - conv.min()
bins = 10
conv_binned = torch.trunc(conv * bins/conv.max().item())
ones = torch.ones_like(conv_binned, requires_grad=True)
zeros = torch.zeros_like(conv_binned, requires_grad=True)
hist = torch.tensor([torch.where(conv_binned == bin, ones, zeros).sum() 
for bin in range(bins)], requires_grad=True)

print(hist.grad_fn)
print(conv_binned.grad_fn)

None
<TruncBackward object at 0x7fa4483a0b38>
I think this is the issue. Any ideas how to fix this? Thanks.