Error in cross entropy loss calculation 160/train.py

[hedongyan]

When I run the PriorFittingCustomPrior.ipynb, there is an error in line 160 of train.py file.

RuntimeError: weight tensor should be defined either for all 2 classes or no classes but got weight tensor of shape: [10]

python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs) 1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks 1050 or _global_forward_hooks or _global_forward_pre_hooks): -> 1051 return forward_call(*input, **kwargs) 1052 # Do not call functions when jit is used 1053 full_backward_hooks, non_full_backward_hooks

~/.conda/envs/tabpfn/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target) 1119 def forward(self, input: Tensor, target: Tensor) -> Tensor: 1120 return F.cross_entropy(input, target, weight=self.weight, -> 1121 ignore_index=self.ignore_index, reduction=self.reduction)

It seems that there is an error in the input of the code line 160 "elif isinstance(criterion, nn.CrossEntropyLoss): losses = criterion(output.reshape(-1, n_out), targets.to(device).long().flatten())".