The present disclosure relates to a method for adapting a computation graph of a machine learning model. The method comprises partitioning the computation graph at an edge between a first node and a second node, wherein the edge is a control dependency edge. The method further comprises inserting a proxy node, a send node, and a receive node between the first node and the second node. The proxy node is connected to the first node, the send node is configured to receive output data from the proxy node, and the receive node is configured to receive the output data from the send node.