title: Training a model with multiple learning rate in PyTorch
body: <p>I am new to PyTorch and getting used to some concepts. </p>
  
  <p>I need to train a Neural Network. For optimization, I need to use Adam optimizer with 4 different <code>learning rates = [2e-5, 3e-5, 4e-5, 5e-5]</code></p>
  
  <p>The optimizer function is defined as below</p>
  
  <pre><code>def optimizer(no_decay = ['bias', 'gamma', 'beta'], lr=2e-5):
      param_optimizer = list(model.named_parameters())
      optimizer_grouped_parameters = [
      {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
       'weight_decay_rate': 0.01},
      {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
       'weight_decay_rate': 0.0}
       ]
       # This variable contains all of the hyperparemeter information our training loop needs
       optimizer = BertAdam(optimizer_grouped_parameters, lr, warmup=.1)
       return optimizer
  </code></pre>
  
  <p>How do I make sure the optimizer uses my specified set of learning rate and returns the best model?</p>
  
  <p>During training, we use the optimizer like below where I don't see a way to tell it to try different learning rate</p>
  
  <pre><code>def model_train():
      #other code
      # clear out the gradient
      optimizer.zero_grad()
      # Forward pass
      loss = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
      train_loss_set.append(loss.item())
      # Backward pass
      loss.backward()
      # Update parameters and take a step using the computed gradient
      optimizer.step()
  </code></pre>
  
  <p>I know that the <code>optimizer.step()</code> internally steps through to optimize the gradient. But how do I make sure the optimizer tries my specified set of learning rate and returns the best model to me?</p>
  
  <p>Please suggest.</p>