I’m not sure about the TimestampCaptureCallback approach, but in any case you could add manual synchronizations via:
auto stream = at::cuda::getCurrentCUDAStream(); // or at::cuda::getDefaultCUDAStream())
AT_CUDA_CHECK(cudaStreamSynchronize(stream));

to make sure your timings are correct. Alternatively you should also be able to use cudaDeviceSynchronize directly:
cudaError_t err = cudaDeviceSynchronize();
bool isEQ = err == cudaSuccess;
ASSERT_TRUE(isEQ);