From e8e00609550deb34d62ce0b163193e461d8c5206 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 23 May 2024 10:35:55 -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 124da6f..d183e8f 100644 --- a/norch/nn/loss.py +++ b/norch/nn/loss.py @@ -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)