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[Stream]) –

Return type:

Generator[None, None, None]