src/autograd/ops/reduce.rs:17 and :32:
vec![Tensor::full(&self.shape, grad.item()).to(grad.device())]
grad.item() downloads from the GPU, the full tensor is then built on the host and
uploaded again. That's every sum() and mean() backward, so every loss, doing a
synchronising round trip plus an allocation the size of the whole input.
CrossEntropyBackward (src/autograd/ops/loss.rs:20) and
WeightedCrossEntropyBackward do the same thing.
fastnn_cuda_fill is already declared in ffi.rs and implemented in kernels.cu
and never called. Either that, or a mul_scalar against a ones buffer, keeps this
on device.
src/autograd/ops/reduce.rs:17and:32:grad.item()downloads from the GPU, the full tensor is then built on the host anduploaded again. That's every
sum()andmean()backward, so every loss, doing asynchronising round trip plus an allocation the size of the whole input.
CrossEntropyBackward(src/autograd/ops/loss.rs:20) andWeightedCrossEntropyBackwarddo the same thing.fastnn_cuda_fillis already declared inffi.rsand implemented inkernels.cuand never called. Either that, or a
mul_scalaragainst a ones buffer, keeps thison device.