human
Write me a code to 3 given the following:
3. Training the model
Now, we can define the model and train it. First we will initilize the weights and the biases to be Gaussian(0,0.01), and then use stochastic gradient descent for optimization. To warm start the training, one may also initilize with pre-trainined parameters arg_params using init=mx.initializer.Load(arg_params). 
To save intermediate results, we can optionally use epoch_end_callback = mx.callback.do_checkpoint(model_prefix, 1) which saves the parameters to the path given by model_prefix, and with period every $1$ epoch. To assess the performance, we output $-\mathcal{L}$(minus ELBO) after each epoch, with the command eval_metric = 'Loss' which is defined above. We will also plot the training loss for mini batches by accessing the log and saving it to a list, and then parsing it to the argument batch_end_callback."
Please write code following the instructions in jupyter notebook style.
assistant
#Here is code to 3.

# set up the log
nd_iter.reset()
logging.getLogger().setLevel(logging.DEBUG)  

# define function to trave back training loss
def log_to_list(period, lst):
    def _callback(param):
        """The checkpoint function."""
        if param.nbatch % period == 0:
            name, value = param.eval_metric.get()
            lst.append(value)
    return _callback

# define the model
model = mx.mod.Module(
    symbol = output ,
    data_names=['data'],
    label_names = ['loss_label'])

# training the model, save training loss as a list.
training_loss=list()

# initilize the parameters for training using Normal.
init = mx.init.Normal(0.01)
model.fit(nd_iter,  # train data
          initializer=init,
          # if eval_data is supplied, test loss will also be reported
          # eval_data = nd_iter_test,
          optimizer='sgd',  # use SGD to train
          optimizer_params={'learning_rate':1e-3,'wd':1e-2},  
          # save parameters for each epoch if model_prefix is supplied
          epoch_end_callback = None if model_prefix==None else mx.callback.do_checkpoint(model_prefix, 1),
          batch_end_callback = log_to_list(N/batch_size,training_loss), 
          num_epoch=100,
          eval_metric = 'Loss')

ELBO = [-training_loss[i] for i in range(len(training_loss))]
plt.plot(ELBO)
plt.ylabel('ELBO');plt.xlabel('epoch');plt.title("training curve for mini batches")
plt.show()