Training a GNN on Pykeen integrated with neo4j

Hi all,

Trying to train a GNN using the Pykeen library, integrated with neo4j to extract the data and collate them into triples. This is the code that I'm using:

training = pipeline(random_seed=random.randint(0,99999),
training=train_data,
testing=test_data,
model=model,
loss=loss,
negative_sampler=negative_sampler,
model_kwargs=dict(embedding_dim=embedding_dim),
optimizer_kwargs=dict(lr=lr),
training_kwargs=dict(num_epochs=epochs, use_tqdm_batch=False, batch_size = train_data.num_triples),
evaluation_kwargs=dict(use_tqdm=False),
evaluation_relation_whitelist = [relation_label])

Currently the model takes about 10 seconds to train on 600 records. I just wanted to ask the community if anyone knows a way to speed this up?