def train_loop_fn(loader, epoch):
    train_steps = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
    tracker = xm.RateTracker()
    total_samples = 0
    rate_list = []
    model.train()
    for (step, (data, target)) in enumerate(loader):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        xm.optimizer_step(optimizer)
        tracker.add(FLAGS.batch_size)
        total_samples += data.size()[0]
        if lr_scheduler:
            lr_scheduler.step()
        if ((step % FLAGS.log_steps) == 0):
            xm.add_step_closure(_train_update, args=(device, step, loss, tracker, epoch, writer))
            test_utils.write_to_summary(writer, step, dict_to_write={'Rate_step': tracker.rate()}, write_xla_metrics=False)
        if (step == train_steps):
            break
    reduced_global = xm.mesh_reduce('reduced_global', tracker.global_rate(), np.mean)
    return (total_samples, reduced_global)
def train_loop_fn(loader, epoch):
    train_steps = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
    tracker = xm.RateTracker()
    total_samples = 0
    rate_list = []
    model.train()
    for (step, (data, target)) in enumerate(loader):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        xm.optimizer_step(optimizer)
        tracker.add(FLAGS.batch_size)
        total_samples += data.size()[0]
        if lr_scheduler:
            lr_scheduler.step()
        if ((step % FLAGS.log_steps) == 0):
            xm.add_step_closure(_train_update, args=(device, step, loss, tracker, epoch, writer))
            test_utils.write_to_summary(writer, step, dict_to_write={'Rate_step': tracker.rate()}, write_xla_metrics=False)
        if (step == train_steps):
            break
    reduced_global = xm.mesh_reduce('reduced_global', tracker.global_rate(), np.mean)
    return (total_samples, reduced_global)