Hello,
If you look at the CrossEntropyLoss documentation, there is a parameter weight that you can specify to give more weight to certain classes. For example,
    weights = torch.tensor([0.1, 0.1, 0.1, 0.3, 0.3, 0.2], dtype=torch.float32)
    criterion = nn.CrossEntropyLoss(weight=weights)

which would give more weight to classes 3, 4, 5.
Hope this helps!