Hi! There’s no way to make autograd aware of the fact that your model is spread across multiple nodes (we have a distributed mode like this, but it needs some work to be completed, and it will take some time to get there).
We don’t have an official example of model parallelism, however it would be sth along the lines of this (the code is incomplete and might contain minor errors, but you get the idea):
import torch.distributed as dist

dist.init_process_group(...)
if dist.rank == 0:
  model = ... # create first part
else:
  model = ... # create second part

if dist.rank == 0:
  for data, labels in data_loader:
    optimizer.zero_grad()
    data = Variable(data)
    output = model(data)
    dist.send(output.data)
    # Now, we're done with our forward part. Wait for gradients
    dist.recv(grad_output)
    output.backward(grad_output)
    optimizer.step() # Do optimization on first part
else:
  for data, labels in data_loader:
    optimizer.zero_grad()
    labels = Variable(labels)
    dist.recv(first_part_output) # Wait for output of first part of the model

    # Compute the final part and the derivatives of this part
    input = Variable(first_part_output, requires_grad=True)
    output = model(input)
    loss = loss_fn(output, labels)
    loss.backward()
    dist.send(input.grad.data) # Send gradients back to node 0
    optimizer.step() # Do optimization on second part