Having problems with my flexible model - as I'm trying to find the reason why it outputs flat lr. Earlier, when I run "normal" [without the option to apply different numbers of layers etc] nn model it actually run properly, having the descending loss. Here's the code: class FlexibleModel(nn.Module): def __init__(self, in_features, hidden_layers_sizes, output_size=1): super().__init__() self.layers = nn.ModuleList() # Use ModuleList to store layers current_in_features = in_features for h_size in hidden_layers_sizes: self.layers.append(nn.Linear(current_in_features, h_size)) current_in_features = h_size # Output of current layer becomes input for next self.out = nn.Linear(current_in_features, output_size) def forward(self, x): for layer in self.layers: x = F.relu(layer(x)) # Apply ReLU after each hidden layer x = self.out(x) return x Lr and configs: hidden_config_1 = [128, 64] model_1 = FlexibleModel(in_features=780, hidden_layers_sizes=hidden_config_1) print("Model 1 Architecture:") print(model_1) optimizer = torch.optim.Adam(model_1.parameters(), lr=0.001) criterion = nn.MSELoss() Training loop num_epochs = 1000 loss_list1 = [] for epoch in range(num_epochs): model_1.train() # Forward pass y_pred = model_1(X_train) loss = F.mse_loss(y_pred, y_train) # Backward + optimize optimizer.zero_grad() loss.backward() optimizer.step() loss_list1.append(loss.item()) # Validation model_1.eval() with torch.no_grad(): val_pred = model_1(X_val) val_loss = F.mse_loss(val_pred, y_val) Here's the plot: I tried gemini and chat to solve this, but no clue, they hallucinate pretty much everytime. Continue reading...