The size of tensor a (2) must match the size of tensor b (3) while defining cross entropy function

I am using the following snippet for defining cross entropy function

def my_cross_entropy(p, y):
"""
Args:
    p: torch tensor with size of (N, C),
    y (int): torch tensor with size of (N), the values range from 0 to C-1

Return:
    loss: the cross_entropy of predicted values p and target y.
"""
loss = None
loss = -np.sum(y*np.log(p))
return loss/float(p.shape[0])

I am making the following method call

batch_size = 2
x = torch.randn(batch_size, 3, 32, 32)
tmp_tensor = torch.randint(3, (batch_size,))
torch_cross_entropy_out = F.cross_entropy(x[::, ::, 0, 0], tmp_tensor)

However, I am getting the following error

RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at 
non-singleton dimension 0

Can anyone help with this?