| 263 | self._qparams = qparams |
| 264 | |
| 265 | def __dlpack__(self, stream=None): |
| 266 | if stream is not None and not isinstance(stream, int): |
| 267 | raise TypeError("stream must be ``int`` or ``none``") |
| 268 | elif stream is not None and stream != -1: |
| 269 | mdevice, mrank, mstream = self.device.physical_locator |
| 270 | if mdevice == "gpu": |
| 271 | if mstream != stream: |
| 272 | device = "gpu{}:{}".format(mrank, mstream) |
| 273 | self.to(device, _borrow=True) |
| 274 | elif mdevice == "cpu": |
| 275 | device = "cpu{}:{}".format(mrank, mstream) |
| 276 | self.to(device, _borrow=True) |
| 277 | else: |
| 278 | raise ValueError("dlpack not support this device: {}!".format(mdevice)) |
| 279 | return _to_dlpack(self) |
| 280 | |
| 281 | |
| 282 | set_py_tensor_type(Tensor) |