How to return specific nodes with the same degree value in a graph in PyTorch?

Viewed 98

I have a function that takes a graph and finds the degree of each individual node then calculates the number of nodes that have the same degree value:

import torch
from torch_geometric.utils import degree

def fun(self, graph):
        
        n = graph.num_nodes
        d = degree(graph.edge_index[1], n, dtype=torch.long)
        counts = torch.bincount(d)

        return counts

The above function works fine. But I want it to only finds the nodes that have a degree less than 50 and then return the count of nodes (of degree < 50) with the same degree value, so I changed the code into the following one:

def fun(self, graph):
        
        n = graph.num_nodes
        deg = degree(graph.edge_index[1], n, dtype=torch.long)
        less_than_fifty = [i if i < 50 else 0 for i in deg]
        counts = torch.bincount(less_than_fifty)

        return counts

After running it, the following error was raised:

TypeError: bincount(): argument 'input' (position 1) must be Tensor, not list

Therefore, I used tensor instead of list as follows:

def fun(self, graph):
        
        n = graph.num_nodes
        d = degree(graph.edge_index[1], n, dtype=torch.long)
        less_than_fifty = torch.tensor([i if i < 50 else 0 for i in deg])
        counts = torch.bincount(less_than_fifty)

        return counts

But this time another issue occured. I ran the code on Google Colab, due to the last modification (converting the list into tensor), Colab kept crashing. I was using GPU on Colab. I am sure the reason of crashing was the line less_than_fifty = torch.tensor([i if i < 50 else 0 for i in deg]) as whenever I deleted it, there was no crash anymore. My question is how to fix the issues?

0 Answers
Related