AnnaKatharinaD:


The error is due to the fact that torch.softmax(...) function expects a tensor & a dimension along which to compute the softmax. You can either use torch.nn.Softmax(..) which by default computes the softmax across the last layer or call torch.softmax(output, dim=-1)