Added necessary changes to make PR #6 work accordingly