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().
- Parameters
stream (Optional[torch.cuda.streams.Stream]) –
- Return type
Generator[None, None, None]