Based on your description, it seems you are trying to classify your dataset into one of three classes.
You are right about the docs of nn.CrossEntropyLoss, i.e. the target might be multi-dimensional.
However, if you compare the shapes of the input (model output) and the target, you see that the channel dimension is missing (C in the docs).
In your case, your target should have the shape [batch_size] and contain the class indices, i.e. values in the range [0, 2].
Just call y_train_torch = torch.argmax(y_train_torch) and you should be fine.
The multi-dimensional use case is useful for e.g. segmentation tasks, where each pixel belongs to one particular class.