code: def train(net, X_train, y_train, epochs, verbose_epoch, learning_rate,
            weight_decay, batch_size):
      
      dataset_train = gluon.data.ArrayDataset(X_train, y_train)
      data_iter_train = gluon.data.DataLoader(dataset_train, batch_size,
                                              shuffle=True)
      trainer = gluon.Trainer(net.collect_params(), 'adam',
                              {'learning_rate': learning_rate,
                               'wd': weight_decay})
      net.initialize(force_reinit=True)
      for epoch in range(epochs):
          for data, label in data_iter_train:
              with autograd.record():
                  output = net(data)
                  loss = square_loss(output, label)
              loss.backward()
              trainer.step(batch_size)
              avg_loss = get_rmse_log(net, X_train, y_train)
          if epoch > verbose_epoch:
              print("Epoch %d, train loss: %f" % (epoch, avg_loss))
      return avg_loss
docs: Trains the model.