neighbors memory update error

[JakubReha]

Hi,
I get the following error when running on a CPU and setting the parameter deliver_to to "neighbors" for the TGN model. Running it on a GPU gives a similar CUDA error. All other parameters are default and it happens across different datasets.

Traceback (most recent call last):
  File "main.py", line 140, in <module>
    test_score, val_score, train_score, res_args, res_conf = link_prediction_single(args, conf)
  File "/home/TGL/train_link.py", line 316, in link_prediction_single
    group_indexes, loss_fun, df[:train_edge_end], rng, device)
  File "/home/TGL/train_link.py", line 77, in train
    mailbox.update_mailbox(model.memory_updater.last_updated_nid, model.memory_updater.last_updated_memory, root_nodes, ts, mem_edge_feats, block)
  File "/home/TGL/memorys.py", line 128, in update_mailbox
    mail = torch.cat([mail, mail[block.edges()[0].long()]], dim=0)
IndexError: index 452 is out of bounds for dimension 0 with size 400
Traceback (most recent call last):
  File "main.py", line 140, in <module>
    test_score, val_score, train_score, res_args, res_conf = link_prediction_single(args, conf)
  File "/home/TGL/train_link.py", line 316, in link_prediction_single
    group_indexes, loss_fun, df[:train_edge_end], rng, device)
  File "/home/TGL/train_link.py", line 77, in train
    mailbox.update_mailbox(model.memory_updater.last_updated_nid, model.memory_updater.last_updated_memory, root_nodes, ts, mem_edge_feats, block)
  File "/home/TGL/memorys.py", line 128, in update_mailbox
    mail = torch.cat([mail, mail[block.edges()[0].long()]], dim=0)
IndexError: index 452 is out of bounds for dimension 0 with size 400