Should we set non_blocking to True?

[ptrblck]

You could use the current stream or the custom stream (if you are using one) and synchronize it as seen e.g. in this example.