code: def use_stream(
      stream: Optional[AbstractStream],
  ) -> Generator[None, None, None]:
      """:func:`torch.cuda.stream` for either CPU or CUDA stream."""
      if not stream:
          yield
          return
  
      if not is_cuda(stream):
          yield
          return
  
      with torch.cuda.stream(as_cuda(stream)):
          yield
  
file_path: fairscale/fairscale/nn/pipe/stream.py
imports: from contextlib import contextmanager
  from typing import Generator, List, Optional, Union, cast
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: use_stream(stream) without return types