degree.py 193 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
import torch


4
def node_degree(row, num_nodes, out=None):
rusty1s's avatar
rusty1s committed
5
    zero = torch.zeros(num_nodes, out=out)
6
7
    one = torch.ones(row.size(0), out=zero.new())
    return zero.scatter_add_(0, row, one)