Context manager to enter compute stream. Equivalent to calling torch.cuda.stream(rc.compute_stream), but also works correctly when cuda is not supported (and compute_stream is None).