def METHOD_NAME(engine, _):
    data = torch.rand(4, 100, device=device)
    model.train()
    data = data.to(device)
    optimizer.zero_grad()
    output = model(data)
    loss = output.sum()
    loss.backward()
    xm.optimizer_step(optimizer, barrier=True)
    return loss.item()