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

1

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.

pfofilr
pfofilr1206×391 14.9 KB

2

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.