Fixes #366.
Moves code to avoid errors.
Merge pull request #367 from LearningLocker/issue/366

Fixes #366.