57c71199259bc937e5a58770cd3733fd65e0ef84,torch_geometric/utils/to_dense_batch.py,,to_dense_batch,#Any#Any#Any#,5

Before Change

    index = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
    index = (index - cum_nodes[batch]) + (batch * max_num_nodes)

    size = [batch_size * max_num_nodes] + list(x.size())[1:]
    batch_x = x.new_full(size, fill_value)
    batch_x[index] = x
    size = [batch_size, max_num_nodes] + list(x.size())[1:]
    batch_x = batch_x.view(size)

After Change

    :rtype: (:class:`Tensor`, :class:`LongTensor`)
    
    batch_size = batch[-1].item() + 1
    num_nodes = scatter_("add", batch.new_ones(x.size(0)), batch, batch_size)
    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
    max_num_nodes = num_nodes.max().item()

Instances