am2.scala

def forwardHandle: Receive = {
  case Forward(x) =>                         // accept forward messages
    val result = module(x)                   // calculate the result
    wire.next ! Forward(result)              // forward the result 
    context become backwardHandle(x, result) // change state to backward
}

def backwardHandle(input: Variable, output: Variable): Receive = {
  case Backward(g) =>                       // accept backward messages
    optimizer.zeroGrad()                    // reset gradients to 0
    output.backward(g)                      // calculate local gradients
    wire.prev ! Backward(input.grad)        // send local gradients
    optimizer.step()                        // update the weights
    context become forwardHandle            // change state to forward
}