Enter the context and set the current stream.
(self)
| 55 | self.stream = stream |
| 56 | |
| 57 | def __enter__(self) -> StreamContext: |
| 58 | """Enter the context and set the current stream.""" |
| 59 | self.prev_stream = core._env_set_current_stream( |
| 60 | self.device_type, self.device_id, self.stream |
| 61 | ) |
| 62 | return self |
| 63 | |
| 64 | def __exit__(self, *args: Any) -> None: |
| 65 | """Exit the context and restore the previous stream.""" |