[TH] code cleanup from PR review