fix autograd was not running on device

2024-05-23 12:07:17 -03:00 · 2024-05-23 12:07:17 -03:00 · 27feb8a366
commit 27feb8a366
parent 472f3189c2
1 changed files with 1 additions and 1 deletions
--- a/norch/tensor.py
+++ b/norch/tensor.py
@ -205,7 +205,7 @@ class Tensor:
                
        if gradient is None:
            if self.shape == [1]:
-                gradient = Tensor([1])
+                gradient = Tensor([1]).to(self.device)
            else:
                raise RuntimeError("Gradient argument must be specified for non-scalar tensors.")