torch.cuda.set_stream
torch.cuda.set_stream(stream)
Sets the current stream.This is a wrapper API to set the stream.
      Usage of this function is discouraged in favor of the "stream"
      context manager.
Parameters:
      stream (Stream) -- selected stream. This function is a no-
      op if this argument is "None".