Fix tree import