Fix for commit 87dcee6a It uses the wrong variable and only covers some use cases. This change covers all use cases.