nn.CrossEntropyLoss expects model outputs with a class dimension as [batch_size, nb_classes, *additional_dims], while the target should not contain this class dimension but instead [batch_size, *additional_dims] and its values should contain the class indices in the range [0, nb_classes-1] as described in the docs.
You are not supposed to set a batch size for any layer or criterion in PyTorch and I guess your target might be one-hot encoded and has thus the additional unwanted dimension.
If that’s the case, use target = torch.argmax(target, dim=1) to create the target tensor in its expected shape.
PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier.