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:41:52 -03:00
parent e8e0060955
commit 472f3189c2
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

@ -67,7 +67,7 @@ class CrossEntropyLoss(Loss):
# target -> Ground truth class indices:
num_classes = input.shape[1]
target = norch.one_hot_encode(target, num_classes)
target = norch.one_hot_encode(target, num_classes).to(target.device)
batch_size = input.shape[0]
logits = norch.softmax(input, dim=1)