Revert "Removed some unnecessary code"

This reverts commit 6261fbf3171c1885fabadf2424b25a87149213a8.