Bug fix for segmentation