diff --git a/XilinxProcessorIPLib/drivers/dp/src/xdp_mst.c b/XilinxProcessorIPLib/drivers/dp/src/xdp_mst.c
index ad1957c3..328279e5 100644
--- a/XilinxProcessorIPLib/drivers/dp/src/xdp_mst.c
+++ b/XilinxProcessorIPLib/drivers/dp/src/xdp_mst.c
@@ -2410,6 +2410,9 @@ void XDp_RxMstExposePort(XDp *InstancePtr, u8 PortNum, u8 Expose)
 	else if (InstancePtr->RxInstance.Topology.LinkAddressInfo.NumPorts) {
 		InstancePtr->RxInstance.Topology.LinkAddressInfo.NumPorts--;
 	}
+
+	XDp_WriteReg(InstancePtr->Config.BaseAddr, XDP_RX_SINK_COUNT,
+		InstancePtr->RxInstance.Topology.LinkAddressInfo.NumPorts);
 }
 
 /******************************************************************************/
@@ -2500,11 +2503,6 @@ void XDp_RxMstSetInputPort(XDp *InstancePtr, u8 PortNum,
 		Branch->Guid[3] = PortOverride->Guid[3];
 	}
 
-	/* Make sure that the branch device is accounted for. */
-	if (InstancePtr->RxInstance.Topology.LinkAddressInfo.NumPorts == 0) {
-		InstancePtr->RxInstance.Topology.LinkAddressInfo.NumPorts = 1;
-	}
-
 	XDp_RxMstExposePort(InstancePtr, PortNum, 1);
 }