Is the calculation of optimizer.step() on CPU?

ddmm2020 · December 17, 2021, 3:29pm

When I do profile on my train code, I found that there would be aten::copy_ and aten::to operations before optimizer.step().So I wonder if optimizer.step() is calculated on cpus.

ptrblck · December 18, 2021, 1:25am

If your model’s parameters are on the GPU the step() operation would also use the GPU to update the parameters. I don’t know which iteration you are profiling, but guess it might be the first one and you might see the initialization of some buffers here. The update can be seen here.