Hi @richard, thanks for your kind help. Finally I changed my code in some other way but got another error saying 'the derivative for 'weight' is not implemented', so it seems impossible to optimize the weights in CrossEntropyLoss function, and on the official website, it is said manual rescaling weight. Maybe I need to use some log function from torch to implement the cross entropy loss by myself.