Fixed bug when running on a CPU device
Merge pull request #535 from peastman/master

Fixed bug when running on a CPU device