aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix cross entropy loss send to device

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 10:35:55 -03:00
parent 8b19ac38c9
commit e8e0060955
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
norch/nn/loss.py
View file

@ -44,7 +44,7 @@ class CrossEntropyLoss(Loss):
if input.ndim == 1:
if target.numel == 1:
num_classes = input.shape[0]
target = norch.one_hot_encode(target, num_classes)
target = norch.one_hot_encode(target, num_classes).to(target.device)
logits = norch.softmax(input, dim=0)
target = target.reshape(logits.shape)