Returns a context manager that specifies the default gpu stream to use. Args: stream_idx: The index of gpu stream to use in the context. Yields: A context manager that specifies the default gpu stream to use for newly created ops. Raise: RuntimeError: If device
(self, stream_idx=0)
| 4422 | |
| 4423 | @tf_contextlib.contextmanager |
| 4424 | def stream(self, stream_idx=0): |
| 4425 | """Returns a context manager that specifies the default gpu stream to use. |
| 4426 | |
| 4427 | Args: |
| 4428 | stream_idx: The index of gpu stream to use in the context. |
| 4429 | |
| 4430 | Yields: |
| 4431 | A context manager that specifies the default gpu stream to use for newly |
| 4432 | created ops. |
| 4433 | |
| 4434 | Raise: |
| 4435 | RuntimeError: If device scopes are not properly nested. |
| 4436 | """ |
| 4437 | self._gpu_stream_stack.push_obj(stream_idx, offset=2) |
| 4438 | old_top_of_stack = self._gpu_stream_stack.peek_top_obj() |
| 4439 | try: |
| 4440 | yield |
| 4441 | finally: |
| 4442 | new_top_of_stack = self._gpu_stream_stack.peek_top_obj() |
| 4443 | if old_top_of_stack is not new_top_of_stack: |
| 4444 | raise RuntimeError("Exiting stream scope without proper scope nesting.") |
| 4445 | self._gpu_stream_stack.pop_obj() |
| 4446 | |
| 4447 | def _apply_gpu_stream(self, op): |
| 4448 | """Applies the current gpu stream stack to the given operation.""" |