ptrblck:

This should work:


Thank you very much, it works for me. But I have another question. After I add cudaStreamSynchronize in my code, the time consumption of tensor.cpu() does reduced. However, I found out that cudaStreamSynchronize  itself cost lots of time, like 70ms which is almost the same time tensor.cpu() cost before I add cudaStreamSynchronize to the code. So, whether I use cudaStreamSynchronize  or not, this extra time consumption is inevitable, is that right? If it is right, is there any way to reduce such time? Like reduce image number of the input tensor?