Updated the example code

Included missing imports