Model architecture does not have H*W*Q size tensors for (a,b) probability

[tathagatv]

As mentioned in the paper, the loss function uses Z, Z^ which are of shape H*W*Q. However the model architecture computes the (a,b) probability distribution in a tensor of shape H/4 * W/4 * Q. How are we computing Z^ then? Is it not supposed to be predicted by the model and we use it for calculation of the loss?

[simonebonato]

We are also having problems because of this... If anyone could enlighten us it would be appreciated.