Handle tensor transfers between devices in the backprop.

2025-06-20 04:00:28 +00:00 · 2023-06-23 08:55:34 +01:00
parent 3f79d81b6f
commit 3e7cb18d7f
2 changed files with 17 additions and 2 deletions
--- a/src/op.rs
+++ b/src/op.rs
@ -17,6 +17,7 @@ pub(crate) enum Op {
    Neg(Tensor),
    Sqr(Tensor),
    Sqrt(Tensor),
+    ToDevice(Tensor),
    Transpose(Tensor, usize, usize),
    // TODO: Support for custom ops.
 }