Message-passing cpu usage and multi-processing

Questions
#1

I try to do some preprocessing on a graph. But an error named DefaultCPUAllocator: can't allocate memory was raised while doing a matrix multipication inside the message function.
I use top command and find that %CPU is increasing to about 500 and %MEM is about 0.5 until the error is raised.

def message_func(edges): 
    ##  shape [15335484, 1]
    is_same = edges.data['is_same ']
    sim= edges.data['sim']
    edge_sims = torch.where(is_same==1, sim, -sim)
    return {'edge_sims': edge_sims}

Then I write a similar code purely by pytorch, and it runs pretty fast.

s= torch.arange(15335484)
x  =  torch.arange(15335484)
y = torch.where(s> 100000, 5* x, -x)

So why this errror happens and how could I do to overcome this problem?

#3

Sorry for my stupid question. I should write my multiplicaiton code in the reduce function.