Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# converting list of class weights to a tensor
weights= torch.tensor(class_weights,dtype=torch.float)
# push to GPU
weights = weights.to(device)
# define the loss function
cross_entropy = nn.NLLLoss(weight=weights)
# number of training epochs
epochs = 10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment