NeighborSampler duplicates batch size - implementation question

[Guykor]

Hello :)

In the NeighborSampler in the sample method, the batch initialize with the following code:

# sample nodes to form positive edges. we will try to predict these edges
row, col, e_id = self.adj_t_sample.coo()
# NOTE: only does self loops when no edges in the current partition of the dataset
target_batch = random_walk(row, col, source_batch, walk_length=1, coalesced=False)[:, 1]
batch = torch.cat([source_batch, target_batch], dim=0)


batch_size: int = len(batch)
        adjs = []
        n_id = batch
        for size in self.sizes:
            adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False)
.....

For an isolated node u, the random_walk returns u in a walk of length 1, but for other nodes the walk will return a random neighbor.
Either way, the new batch, which now consist of duplicate amount of nodes, is then fed to the adj sampler. Hereby, sampling twice from the same nodes, or sampling neighbors from nodes that didn't appeared in the original batch.
I don't fully understand how this achieves self-loop for isolated nodes, wouldn't it make more sense just to add the missing edges to obtain self-loop?

Thank you :)