Record the stream that is using this graph. This method only supports the PyTorch backend and requires graphs on the GPU. Parameters ---------- stream : torch.cuda.Stream The stream that is using this graph. Returns ------- DGLGra
(self, stream)
| 5872 | return self._graph.is_pinned() |
| 5873 | |
| 5874 | def record_stream(self, stream): |
| 5875 | """Record the stream that is using this graph. |
| 5876 | This method only supports the PyTorch backend and requires graphs on the GPU. |
| 5877 | |
| 5878 | Parameters |
| 5879 | ---------- |
| 5880 | stream : torch.cuda.Stream |
| 5881 | The stream that is using this graph. |
| 5882 | |
| 5883 | Returns |
| 5884 | ------- |
| 5885 | DGLGraph |
| 5886 | self. |
| 5887 | """ |
| 5888 | if F.get_preferred_backend() != "pytorch": |
| 5889 | raise DGLError("record_stream only support the PyTorch backend.") |
| 5890 | if F.device_type(self.device) != "cuda": |
| 5891 | raise DGLError("The graph must be on GPU to be recorded.") |
| 5892 | self._graph.record_stream(stream) |
| 5893 | for frame in itertools.chain(self._node_frames, self._edge_frames): |
| 5894 | for col in frame._columns.values(): |
| 5895 | col.record_stream(stream) |
| 5896 | |
| 5897 | return self |
| 5898 | |
| 5899 | def clone(self): |
| 5900 | """Return a heterograph object that is a clone of current graph. |