where predictions is [32, 7], YTrain represents the true labels (0,1,2,…), and varTrain should have the same size of predictions. The code only works at the first epoch, then I get the following error:

Var: (N, *) or (*), same shape as the input, or same shape as the input but with one dimension equal to 1, or same shape as the input but with one fewer dimension (to allow for broadcasting)

Check the shape of varTrain and make sure it meets these requirements.

Actually, the dimension of varTrain is [32, 7] as the dimension of predictions, indeed, the code seems to work until the first epoch. After that I get the error and I cannot understand the reason. Any suggestion?