[Fix] Minor edge case fix