I have a GNN where I am using the standard node level features and graph level features, right now I have my edge_attr(edge weights) as constant or they just depend upon how far the nodes are separated (I have x,y,z co-ordinates for nodes in my problem use case). Now I want those edge Attributes to be trainable. How do I do that? I tried defining these in the init function but since every batch of graph has different number of edges, the size is variable, One thing I can do is, define a large enough tensor in the init function of size (some large number, 1) and then use it’s subtensor of size (edge_attribute,1) in forward function. Is there a better way for this?