def train_step(input, target, model, loss_fn, optimizer, **unused):
    """Do forward, backward and parameter update."""
    model.train()
    output = model(input)
    loss = loss_fn(output, target)
    optimizer.backward(loss)
    optimizer.step()
Do forward, backward and parameter update.
edgelm/tests/distributed/test_bmuf.py