Fix the sample code

Fixes #4