Hi，I guess it is probably your large number of .cuda() operations during the for loop that makes it quite slow. How about trying to initialize variable “size” out of the loop and assign it in the loop?