Two possible bugs in synthetic dataset code

[GAldegheri]

Hi, I was trying to reproduce the synthetic graph experiments and I found what look like two bugs - both in data.py:
https://github.com/jrwnter/pigvae/blob/c3dfcef252f2bf7d34ee4c8dca2ca5a605fa894b/pigvae/synthetic_graphs/data.py#L159C13-L169
When dm is converted to .long() on line 164, it seems like infinite entries in dm (indicating two unconnected nodes) turn into very large negative numbers. This causes the subsequent clamping to turn them into zeros. The result is that both unconnected nodes and self-connections are coded as the first entry in the one-hot vectors. I don't know how big of a problem this is, but it's relatively straightforward to fix by moving the conversion to 'long' after the clamping:

dm = torch.from_numpy(floyd_warshall_numpy(graph))
dm = torch.clamp(dm, 0, 5).unsqueeze(-1).long()

Another issue is that after num_nodes is reassigned on line 166, it always corresponds to the maximum number of nodes. This results in the mask always containing True exclusively. This is also straightforward to fix like this (i.e. without reassigning num_nodes):

dm = torch.zeros((max_num_nodes, max_num_nodes, 6)).type_as(dm).scatter_(2, dm, 1).float()
edge_features.append(dm)
mask.append((torch.arange(max_num_nodes) < num_nodes).unsqueeze(0))

Please let me know if I misunderstood something and these are not actually bugs!