my code doesn't work

[businessabstraction]

Hi (help!!),
I have a randomly created bipartite graph. Input nodes have hidden values. The output nodes are labled with maximum of the hidden values. IN_DEGREES are in-degrees of the target nodes(all the same)
In my model, I try to find the hidden values, and the

 in2out = []
    max_vals = []
    for ind in range(TRAIN_OUTPUTS+TEST_OUTPUTS):
    in_indexes =  random.sample(range(len(data['In'].x)), IN_DEGREE)
          highest_val  = max([data['In'].x[hi] for hi in in_indexes])
          max_vals.append(highest_val)
          assert ind + 1 == len(max_vals)
          for hi in in_indexes:
            in2out.append([hi, ind])
    return max_vals, in2out
def makeData():
    data = HeteroData()
    data['In'].x = torch.tensor([random.randint(MIN_VAL, MAX_VAL) for i in range(NUM_INPUTS)])
    assert len(data['In'].x) == NUM_INPUTS
    max_vals, in2out = makeEdges(data)
    
    data['Out'].maximums = torch.tensor(max_vals, dtype=torch.float32)
    data['In', 'Out'].edge_index = torch.tensor(in2out, dtype=torch.int64).t().contiguous()
    data['In', 'Out'].sz = (len(data['In'].x), len(max_vals))
    data['Out'].train_mask = torch.tensor([True]*TRAIN_OUTPUTS + [False]*TEST_OUTPUTS)
    data['Out'].test_mask = torch.logical_not(data['Out'].train_mask)
    return data

I created the most basic model that should find the hidden values using the training nodes, and check using the test node. One would think it should be able to find the values precisely:

class In2Out(pyg.nn.MessagePassing):
    '''  straight through message passing  with max aggregation '''
    def __init__(self):
        super().__init__(aggr="max", flow="source_to_target")
        
    def reset_parameters(self):
        pass
        
    def forward(self, edge_index, size, x, report=False):
        out = self.propagate(edge_index, size=size, x=x, report=report)
        return out
    
    def message(self, x_j, report = False):
        return x_j

class maxSpeedGNN(torch.nn.Module):
    """ Network to guess the hidden values and the maximums """ 
    def __init__(self, data, report=False):
        super(maxSpeedGNN, self).__init__()
        self.i2o = In2Out()        
        input_num = data['In'].x.size()[0]
        self.input_embedding = torch.nn.Parameter(torch.zeros(input_num, INPUT_DIM))
        torch.nn.init.uniform_(self.input_embedding)          
        return
        
    def forward(self, data, report=False):        
        outputs = self.i2o(data['In', 'Out'].edge_index, data['In', 'Out'].sz, x=self.input_embedding, report=report)     
        return outputs
    

def runMaxReversing(data, model, epochs = 5, report=False):
    #torch.manual_seed(1)
    torch.autograd.set_detect_anomaly(True)
    EPOCH_LENGTH = 1000
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    lf = nn.MSELoss()

    model.train()

    for epoch in range(epochs):
        for i in range(EPOCH_LENGTH):
            optimizer.zero_grad()  # Clear gradients.
            maximums = model(data, report=report)  # Perform a single forward pass.
            loss = lf(maximums[data['Out'].train_mask], data['Out'].maximums[data['Out'].train_mask])  
            loss.backward()  
            optimizer.step() 
        model.eval()
        test_loss = lf(maximums[data['Out'].test_mask], data['Out'].maximums[data['Out'].test_mask])
        print(f"epoch {epoch}, loss {loss}, test loss {test_loss}")
        model.train() 
    return model   

The loss stops at some non-zero value, depending on in-degrees

Cannot understaned why. Attaching the notebook
mpnmostbasic.zip

[akihironitta]

Haven't looked at your code yet, but in general, I'd recommend starting from any of similar examples available at https://github.com/pyg-team/pytorch_geometric/tree/master/examples.

[businessabstraction]

I looked through examples, I found none that learns source embeddings (torch.nn.Parameter) rather than matrices within the Message Passin paradigm. Any advice on which example can be utilised would be hugely appreciated