Validation of the dataset in GCN

Hello!

I am trying to do the validation on my dataset to obtain optimal number of channels in my neural network. I have the following code:

def train_during_validation():
    for epoch in range (1, 201):
        model.train()
        optimizer.zero_grad()
        out = model(data.x, data.edge_index)
        loss = criterion(out[data.val_mask], data.y[data.val_mask])
        loss.backward()
        optimizer.step()
    return loss

def validation():
    loss_val = np.zeros(50, dtype = float)
    model = GCN(hidden_channels = 1)
    loss_val = train_during_validation()
    print(loss_val)
        
validation()

In the code above I train previously defined model with 16 channels and I obtain a loss of 0.33. But as soon as I start doing validation on hidden_channel (see code below), my loss does not go down (it remains on 1.95). I do not understand why. Can somebody explain?

def train_during_validation(model):
    print(f'Model:{model}')
    for epoch in range (1, 201):  
        model.train()
        optimizer.zero_grad()
        out = model(data.x, data.edge_index)
        loss = criterion(out[data.val_mask], data.y[data.val_mask])
        loss.backward()
        optimizer.step()
    return loss

def validation():
    loss_val = np.zeros(50, dtype = float)
    model = GCN(hidden_channels = 1)
    for i in range (50):
        model = GCN(hidden_channels = i)
        #print(model)
        loss_val[i] = train_during_validation(model)
        print(loss_val[i])
        
validation()

Thank you in advance!

Best regards