To be clear, since you already know how to split the model into two model shards, is your question how to convert the following lines in https://github.com/pytorch/vision/blob/master/references/detection/engine.py:
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

into the code in RPC Framework Tutorial?
        with dist_autograd.context() as context_id:
            outputs = model(inputs)
            dist_autograd.backward(context_id, [loss_fn(outputs, labels)])
            opt.step(context_id)

cc: @mrshenli