Also move feature flag check to ensure it is used in all cases. Change-Id: If42787c5bc9ab824449c70f90fb827cf2da7507f