# Buggy code
def _prepare_cleanup(training: np.ndarray, testing: np.ndarray):
    training_entities = _get_unique(training)
    testing_entities = _get_unique(testing)
    to_move = testing_entities[~np.isin(testing_entities, training_entities)]
    move_id_mask = np.isin(testing[:, [0, 2]], to_move).any(axis=1)
    return training_entities, testing_entities, to_move, move_id_mask


# 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 _prepare_cleanup(training: np.ndarray, testing: np.ndarray) -> np.ndarray:
    to_move_mask = None
    for col in [[0, 2], 1]:
        training_ids, test_ids = [
            np.unique(triples[:, col]) for triples in [training, testing]
        ]
        to_move = test_ids[~np.isin(test_ids, training_ids)]
        this_to_move_mask = np.isin(testing[:, col], to_move)
        if this_to_move_mask.ndim > 1:
            this_to_move_mask = this_to_move_mask.any(axis=1)
        if to_move_mask is None:
            to_move_mask = this_to_move_mask
        else:
            to_move_mask = this_to_move_mask | to_move_mask

    return to_move_mask