In the calculation of the NCE loss, the softmax does not have a dimension to compute the result and by default, PyTorch uses dim=1 with 2D input.
The Loss in the paper highlights that the c_t (context) remains constant, and we 'match' this context to the actual values of z_t. By using dim=1 instead of dim=0 we actually compute the 'match' between a constant z_t and c_ts that are generated by each example in the batch.
The softmax should be performed on the columns of the 8x8 matrix to capture the true loss function defined in the CPC paper.
In the calculation of the NCE loss, the softmax does not have a dimension to compute the result and by default, PyTorch uses dim=1 with 2D input.
The Loss in the paper highlights that the c_t (context) remains constant, and we 'match' this context to the actual values of z_t. By using dim=1 instead of dim=0 we actually compute the 'match' between a constant z_t and c_ts that are generated by each example in the batch.
The softmax should be performed on the columns of the 8x8 matrix to capture the true loss function defined in the CPC paper.