Move looping inside the GPU kernels

Otherwise, we'll have to copy from device to host and back again between each
kernel run. Also, not sequentially launching the kernel many thousands of
times helps.
Also copy 'weight' (in addition to 'out') after kernel run