Fix to parallel forward pass.