def test_loop_fn(loader, epoch):
    test_steps = (testsize // (FLAGS.test_set_batch_size * xm.xrt_world_size()))
    (total_samples, correct) = (0, 0)
    model.eval()
    for (step, (data, target)) in enumerate(loader):
        output = model(data)
        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(target.view_as(pred)).sum()
        total_samples += data.size()[0]
        if ((step % FLAGS.log_steps) == 0):
            xm.add_step_closure(test_utils.print_test_update, args=(device, None, epoch, step))
        if (step == test_steps):
            break
    correct_val = correct.item()
    accuracy_replica = ((100.0 * correct_val) / total_samples)
    accuracy = xm.mesh_reduce('test_accuracy', accuracy_replica, np.mean)
    return (accuracy, accuracy_replica, total_samples)
def test_loop_fn(loader, epoch):
    test_steps = (testsize // (FLAGS.test_set_batch_size * xm.xrt_world_size()))
    (total_samples, correct) = (0, 0)
    model.eval()
    for (step, (data, target)) in enumerate(loader):
        output = model(data)
        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(target.view_as(pred)).sum()
        total_samples += data.size()[0]
        if ((step % FLAGS.log_steps) == 0):
            xm.add_step_closure(test_utils.print_test_update, args=(device, None, epoch, step))
        if (step == test_steps):
            break
    correct_val = correct.item()
    accuracy_replica = ((100.0 * correct_val) / total_samples)
    accuracy = xm.mesh_reduce('test_accuracy', accuracy_replica, np.mean)
    return (accuracy, accuracy_replica, total_samples)