From 472f3189c2c2d75fc87eb2d8c6f0edf7f227d7b2 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 23 May 2024 10:41:52 -0300 Subject: [PATCH] fix cross entropy loss send to device --- norch/nn/loss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/norch/nn/loss.py b/norch/nn/loss.py index d183e8f..6c0abea 100644 --- a/norch/nn/loss.py +++ b/norch/nn/loss.py @@ -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)