Hi,
I am trying to minimize the distance between two feature maps using KL Divergence. The output file is of shape in both the case are : (16, 1, 8, 224, 224)
One map is the prediction on RGB clip and another one is the flipped version of the same clip. I’m getting a NaN loss.
output = model(data, label)  # 16, 1, 8, 224, 224
output = torch.flip(output, [4])
output = torch.mean(output, 2) # 16, 1, 224, 224
flip_op = torch.mean(flip_op, 2)
output = torch.squeeze(output, 1). # 16, 224, 224
flip_op = torch.squeeze(flip_op, 1)
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = consistency_criterion(flip_op.log(), output.detach()).sum(-1).mean()