fix mse bug of using old gradient.data (#504)

Co-authored-by: struggle <blank>
mratsim · Apr 17, 2021 · 2c4a39f · 2c4a39f
1 parent 6118841
commit 2c4a39f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/arraymancer/nn/loss/mean_square_error_loss.nim b/src/arraymancer/nn/loss/mean_square_error_loss.nim
@@ -24,7 +24,7 @@ proc mse_backward_ag[TT](self: MSELoss[TT], payload: Payload[TT]): SmallDiffs[TT
   let gradient = payload.variable.grad
   # Gradient is a tensor of shape 1
   assert gradient.shape == [1]
-  let grad = gradient.data[gradient.offset]
+  let grad = gradient.unsafe_raw_offset[0]
 
   let norm = grad * 2'f32 / gradient.size.float32 # TODO divide by total number of elements or by batch size? https://github.com/pytorch/pytorch/issues/3322
                                                   # See also Stanford course: http://theory.stanford.edu/~tim/s15/l/l15.pdf