Hi,
As mentioned in the error message, you should be able to enable anomaly mode to get a better idea where the error comes from.
In the function you mention, doing max[:, 0, :] = XXX is an inplace operation. So most likely the issue?
You can try to replace this with:
max = max * torch.tensor([[[1920,], [1200,], [200,]],]) (weird brackets to make it a tensor of size [1, 3, 1] to make sure broadcasting does what you want.