Fix for draw_samples. Still needs to be optimized.