Re-added Working Lines

After the previous fix, I forgot to add the lines of code that I
removed.