def train(epoch, ctx, train_data, dev_data):
    if isinstance(ctx, mx.Context):
        ctx = [ctx]
    net.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx[0])
    net.embed.weight.set_data(vocab.embed.as_in_context(ctx[0]))
    train_data.set_context(ctx[0])
    dev_data.set_context(ctx[0])
    trainer = gluon.Trainer(net.collect_params(), optimizer, {'learning_rate': opt.lr, 'wd': opt.wd})
    best_r = (- 1)
    Loss = gluon.loss.KLDivLoss()
    for i in range(epoch):
        train_data.reset()
        num_batches = len(train_data)
        preds = []
        labels = [mx.nd.array(train_data.labels, ctx=ctx[0]).reshape(((- 1), 1))]
        for j in tqdm(range(num_batches), desc='Training epoch {}'.format(i)):
            (l_tree, l_sent, r_tree, r_sent, label) = train_data.next()
            with ag.record():
                z = net(mx.nd, l_sent, r_sent, l_tree, r_tree)
                loss = Loss(z, to_target(label).as_in_context(ctx[0]))
                loss.backward()
                preds.append(z)
            if (((j + 1) % batch_size) == 0):
                trainer.step(batch_size)
        preds = to_score(mx.nd.concat(*preds, dim=0))
        metric.update(preds, labels)
        (names, values) = metric.get()
        metric.reset()
        for (name, acc) in zip(names, values):
            logging.info(('training acc at epoch %d: %s=%f' % (i, name, acc)))
        best_r = test(ctx, dev_data, best_r, num_iter=i)
def train(epoch, ctx, train_data, dev_data):
    if isinstance(ctx, mx.Context):
        ctx = [ctx]
    net.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx[0])
    net.embed.weight.set_data(vocab.embed.as_in_context(ctx[0]))
    train_data.set_context(ctx[0])
    dev_data.set_context(ctx[0])
    trainer = gluon.Trainer(net.collect_params(), optimizer, {'learning_rate': opt.lr, 'wd': opt.wd})
    best_r = (- 1)
    Loss = gluon.loss.KLDivLoss()
    for i in range(epoch):
        train_data.reset()
        num_batches = len(train_data)
        preds = []
        labels = [mx.nd.array(train_data.labels, ctx=ctx[0]).reshape(((- 1), 1))]
        for j in tqdm(range(num_batches), desc='Training epoch {}'.format(i)):
            (l_tree, l_sent, r_tree, r_sent, label) = train_data.next()
            with ag.record():
                z = net(mx.nd, l_sent, r_sent, l_tree, r_tree)
                loss = Loss(z, to_target(label).as_in_context(ctx[0]))
                loss.backward()
                preds.append(z)
            if (((j + 1) % batch_size) == 0):
                trainer.step(batch_size)
        preds = to_score(mx.nd.concat(*preds, dim=0))
        metric.update(preds, labels)
        (names, values) = metric.get()
        metric.reset()
        for (name, acc) in zip(names, values):
            logging.info(('training acc at epoch %d: %s=%f' % (i, name, acc)))
        best_r = test(ctx, dev_data, best_r, num_iter=i)