classImbalance.py

from sklearn.utils.class_weight import compute_class_weight

#compute class weights
classw = compute_class_weight('balanced', np.unique(train_labels), trainlabels)

print("Class Weights:",classweights)

#Output obtained: [0.57743559 3.72848948]

# converting list of class weights to a tensor
weight= torch.tensor(classw , dtype=torch.float)

# push to GPU
weight = weight.to(device)

# define the loss function
crossentropy  = nn.NLLLoss(weight=weight) 

# number of training epochs
epochs = 10