Custom widget bug fix