Update the affine fusion implementation to handle split datasets.
Update the affine fusion implementation to handle split datasets.