def train()

in source/sagemaker/baselines/mlp_fraud_entry_point.py [0:0]

• 18 lines of code
• 3 McCabe index (conditional complexity)

def train(model, trainer, loss, train_data, ctx):
    duration = []
    for epoch in range(args.n_epochs):
        tic = time.time()
        loss_val = 0.

        for features, labels, weight_mask in train_data:
            with autograd.record():
                pred = model(features.as_in_context(ctx))
                l = loss(pred, labels.as_in_context(ctx), mx.nd.expand_dims(weight_mask, 1).as_in_context(ctx)).sum()
            l.backward()
            trainer.step(args.batch_size)

        loss_val += l.asscalar()
        duration.append(time.time() - tic)
        f1 = evaluate(model, train_data, ctx)
        logging.info("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | F1 Score {:.4f} ".format(
            epoch, np.mean(duration), loss_val / n, f1))
    save_model(model)
    return model