Thanks for the code.
Could you add synchronization points before starting and stopping the timer for the forward and backward pass?
at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
AT_CUDA_CHECK(cudaStreamSynchronize(stream));

It would be interesting to see just the time for the actual GPU workload.