# Buggy code
def split(
    self,
    ratios: Union[float, Sequence[float]] = 0.8,
    *,
    random_state: Union[None, int, np.random.RandomState] = None,
    randomize_cleanup: bool = False,
) -> List["TriplesFactory"]:
    """Split a triples factory into a train/test.

    :param ratios: There are three options for this argument. First, a float can be given between 0 and 1.0,
     non-inclusive. The first triples factory will get this ratio and the second will get the rest. Second,
     a list of ratios can be given for which factory in which order should get what ratios as in ``[0.8, 0.1]``.
     The final ratio can be omitted because that can be calculated. Third, all ratios can be explicitly set in
     order such as in ``[0.8, 0.1, 0.1]`` where the sum of all ratios is 1.0.
    :param random_state: The random state used to shuffle and split the triples in this factory.
    :param randomize_cleanup: If true, uses the non-deterministic method for moving triples to the training set.
     This has the advantage that it doesn't necessarily have to move all of them, but it might be slower.

    .. code-block:: python

        ratio = 0.8  # makes a [0.8, 0.2] split
        training_factory, testing_factory = factory.split(ratio)

        ratios = [0.8, 0.1]  # makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)

        ratios = [0.8, 0.1, 0.1]  # also makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)
    """
    n_triples = self.triples.shape[0]

    # Prepare shuffle index
    idx = np.arange(n_triples)
    if random_state is None:
        random_state = np.random.randint(0, 2**32 - 1)
        logger.warning(f"Using random_state={random_state} to split {self}")
    if isinstance(random_state, int):
        random_state = np.random.RandomState(random_state)
    random_state.shuffle(idx)

    # Prepare split index
    if isinstance(ratios, float):
        ratios = [ratios]

    ratio_sum = sum(ratios)
    if ratio_sum == 1.0:
        ratios = ratios[:-1]  # vsplit doesn't take the final number into account.
    elif ratio_sum > 1.0:
        raise ValueError(f"ratios sum to more than 1.0: {ratios} (sum={ratio_sum})")

    split_idxs = [int(split_ratio * n_triples) for split_ratio in ratios]
    # Take cumulative sum so the get separated properly
    split_idxs = np.cumsum(split_idxs)

    # Split triples
    triples_groups = np.vsplit(self.triples[idx], split_idxs)
    logger.info(
        f"split triples to groups of sizes {[triples.shape[0] for triples in triples_groups]}"
    )

    # Make sure that the first element has all the right stuff in it
    triples_groups = _tf_cleanup_all(
        triples_groups, random_state=random_state if randomize_cleanup else None
    )

    # Make new triples factories for each group
    return [
        TriplesFactory(
            triples=triples,
            entity_to_id=deepcopy(self.entity_to_id),
            relation_to_id=deepcopy(self.relation_to_id),
        )
        for triples in triples_groups
    ]


# Traceback (IndexError)
Using random_state=1333753659 to split TriplesFactory(path="/tmp/out")
No random seed is specified. Setting to 2098687070.
No cuda devices were available. The model runs on CPU
Training epochs on cpu:   0%|                                                           | 0/5 [00:00<?, ?epoch/s]INFO:pykeen.training.training_loop:using stopper: <pykeen.stoppers.stopper.NopStopper object at 0x7f42e4842cc0>
Training epochs on cpu:   0%|                                                           | 0/5 [00:00<?, ?epoch/s]
Traceback (most recent call last):
File "test.py", line 8, in <module>
training_triples_factory=training, testing_triples_factory=testing, model="TransH"
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/pipeline.py", line 815, in pipeline
**training_kwargs,
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/training/training_loop.py", line 190, in train
num_workers=num_workers,
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/training/training_loop.py", line 376, in _train
slice_size,
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/training/training_loop.py", line 438, in _forward_pass
slice_size=slice_size,
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/training/slcwa.py", line 101, in _process_batch
positive_scores = self.model.score_hrt(positive_batch)
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/pykeen/models/unimodal/trans_h.py", line 134, in score_hrt
d_r = self.relation_embeddings(hrt_batch[:, 1])
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 114, in forward
self.norm_type, self.scale_grad_by_freq, self.sparse)
File "/home/dobraczka/.local/share/virtualenvs/embedding-transformers-I3i1Obsv/lib/python3.7/site-packages/torch/nn/functional.py", line 1724, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
IndexError: index out of range in self

# Fixed code
def split(
    self,
    ratios: Union[float, Sequence[float]] = 0.8,
    *,
    random_state: Union[None, int, np.random.RandomState] = None,
    randomize_cleanup: bool = False,
) -> List["TriplesFactory"]:
    """Split a triples factory into a train/test.

    :param ratios: There are three options for this argument. First, a float can be given between 0 and 1.0,
     non-inclusive. The first triples factory will get this ratio and the second will get the rest. Second,
     a list of ratios can be given for which factory in which order should get what ratios as in ``[0.8, 0.1]``.
     The final ratio can be omitted because that can be calculated. Third, all ratios can be explicitly set in
     order such as in ``[0.8, 0.1, 0.1]`` where the sum of all ratios is 1.0.
    :param random_state: The random state used to shuffle and split the triples in this factory.
    :param randomize_cleanup: If true, uses the non-deterministic method for moving triples to the training set.
     This has the advantage that it doesn't necessarily have to move all of them, but it might be slower.

    .. code-block:: python

        ratio = 0.8  # makes a [0.8, 0.2] split
        training_factory, testing_factory = factory.split(ratio)

        ratios = [0.8, 0.1]  # makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)

        ratios = [0.8, 0.1, 0.1]  # also makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)
    """
    n_triples = self.triples.shape[0]

    # Prepare shuffle index
    idx = np.arange(n_triples)
    if random_state is None:
        random_state = np.random.randint(0, 2**32 - 1)
        logger.warning(f"Using random_state={random_state} to split {self}")
    if isinstance(random_state, int):
        random_state = np.random.RandomState(random_state)
    random_state.shuffle(idx)

    # Prepare split index
    if isinstance(ratios, float):
        ratios = [ratios]

    ratio_sum = sum(ratios)
    if ratio_sum == 1.0:
        ratios = ratios[:-1]  # vsplit doesn't take the final number into account.
    elif ratio_sum > 1.0:
        raise ValueError(f"ratios sum to more than 1.0: {ratios} (sum={ratio_sum})")

    sizes = [int(split_ratio * n_triples) for split_ratio in ratios]
    # Take cumulative sum so the get separated properly
    split_idxs = np.cumsum(sizes)

    # Split triples
    triples_groups = np.vsplit(self.triples[idx], split_idxs)
    logger.info(
        f"split triples to groups of sizes {[triples.shape[0] for triples in triples_groups]}"
    )

    # Make sure that the first element has all the right stuff in it
    triples_groups = _tf_cleanup_all(
        triples_groups, random_state=random_state if randomize_cleanup else None
    )

    for i, (triples, exp_size, exp_ratio) in enumerate(
        zip(triples_groups, sizes, ratios)
    ):
        actual_size = triples.shape[0]
        actual_ratio = actual_size / exp_size * exp_ratio
        if actual_size != exp_size:
            logger.warning(
                f"Requested ratio[{i}]={exp_ratio:.3f} (equal to size {exp_size}), but got {actual_ratio:.3f} "
                f"(equal to size {actual_size}) to ensure that all entities/relations occur in train."
            )

    # Make new triples factories for each group
    return [
        TriplesFactory(
            triples=triples,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            compact_id=False,
        )
        for triples in triples_groups
    ]