pytorch_pfn_extras.cuda.stream¶
- pytorch_pfn_extras.cuda.stream(stream)¶
Context-manager that selects a given stream.
This context manager also changes the CuPy’s default stream if CuPy is available. When CuPy is not available, the functionality is the same as the PyTorch’s counterpart, torch.cuda.stream().