Mapping NCH output of Conv1D to N,H,NUM_CLASSES

I am working with a 1D convolutional layer which outputs a 3D tensor of dimensions N,C,H. I would like to classify the output of each row H into a class so I would like to map the N,C,H tensor to to a tensor (N, H, NUM_CLASSES) or (N, NUM_CLASSES, H)

One way to do this is to have my final layer be a convolutional layer Conv1d(kernel_size=1, in_channels=C, out_channels=NUM_CLASSES, stride=1, padding=0) which would give my output layer the shape (N, NUM_CLASSES, H).

Is this the right way to go about this problem? My label data is a matrix of dimensions (N, NUM_CLASSES).

I am using CrossEntropyLoss for training and the documentation says input has to be a 2D Tensor of size (minibatch, C). so with my current approach of mapping to (N, NUM_CLASSES, H) it seems like my tensor needs to be reshaped.