Fix code for TDIS-5396