Out of memory in message func

below is my message func for my GCN network i am calculating the matrix according to if the edge is a backward or forward
meaning t*forward matrix +(t-1)*backward matrix where t is 1 for a forward edge and 0 for a backward edge but it takes a very large amount of GPU memory any suggestions on how i can improve the memory utilization in this case?

    def massage_func(self, edges):
    forward = self.forward_edge(edges.src['emb'])
    backward = self.backward_edge(edges.src['emb'])
    relation = th.t(edges.data['edge_relation'].expand(embedding_size, len(edges)))
    u = forward * relation
    v = backward * (1 - relation)
    return {'m': th.relu(u + v)}

i think the problem is the mail box i’m doing ten iterations of send and receive and between iteration he wont clear the memory

I suppose a scalar is enough for relation and you don’t really need expand(...)?

How large is your graph in terms of the number of edges?

relation is a matrix we might be able to make it a vector instead of a matrix ,
we use a batched dgl graph so the amount of edge is around tens of thousands for a batch of 20 .

Yeah, that’s not a lot of edges. Try replacing the matrix by a vector and see if it helps. Just for reference, you may also provide the numbers for GPU memory occupied.