def METHOD_NAME(self, idx: int) -> Stream:
    if self._push_streams[idx] is None:
        assert str(self._device)[:4] == 'cuda'
        self._push_streams[idx] = torch.cuda.Stream(self._device)
    return self._push_streams[idx]