It seems that the problem was with the +=:
When separating this line of code into this:
sub_patch[i:(i + filter_sz), j:(j + filter_sz), :] = sub_patch[i:(i + filter_sz), j:(j + filter_sz), :] +(sub_filt_patch * pred_patch[i,j]).sum(dim=3)

Then it worked!
This must be a bug I guess?