Link prediction KeyError

[katarinasupe]

User reported error on Discord:

Query failed: link_prediction.train: Traceback (most recent call last):
  File "/usr/lib/memgraph/query_modules/link_prediction.py", line 322, in train
    training_results, validation_results = inner_train(
  File "/usr/lib/memgraph/query_modules/mage/link_prediction/link_prediction_util.py", line 420, in inner_train
    probs, labels, loss_output = batch_forward_pass(
  File "/usr/lib/memgraph/query_modules/mage/link_prediction/link_prediction_util.py", line 262, in batch_forward_pass
    pos_score = predictor.forward(pos_graph, outputs, target_relation=target_relation)
  File "/usr/lib/memgraph/query_modules/mage/link_prediction/predictors/DotPredictor.py", line 35, in forward
    g.apply_edges(
  File "/usr/local/lib/python3.9/dist-packages/dgl-0.9.1.post1-py3.9-linux-x86_64.egg/dgl/heterograph.py", line 4458, in apply_edges
    edata = core.invoke_gsddmm(g, func)
  File "/usr/local/lib/python3.9/dist-packages/dgl-0.9.1.post1-py3.9-linux-x86_64.egg/dgl/core.py", line 259, in invoke_gsddmm
    x = alldata[func.lhs][func.lhs_field]
  File "/usr/local/lib/python3.9/dist-packages/dgl-0.9.1.post1-py3.9-linux-x86_64.egg/dgl/view.py", line 73, in getitem
    return self._graph._get_n_repr(self._ntid, self._nodes)[key]
  File "/usr/local/lib/python3.9/dist-packages/dgl-0.9.1.post1-py3.9-linux-x86_64.egg/dgl/frame.py", line 622, in getitem
    return self._columns[name].data
KeyError: <Predictors.NODE_EMBEDDINGS: 'node_embeddings'>

He is performing link prediction on heterogeneous graph. There is only one edge type and edges do not have any properties. It was concluded that this is some kind of a edge case.