Add imports to code example in the README

Description

Add imports to the code examples in the README.

Makes sure the code example is complete and runs as it is.

[ ] I have raised an issue to propose this change (required for new features and bug fixes)

[ ] Bug fix (non-breaking change which fixes an issue)
[ ] New feature (non-breaking change which adds functionality)
[ ] Breaking change (fix or feature that would cause existing functionality to change)
[x] Documentation (update in the documentation)

[ ] I have updated the changelog accordingly (required).
[ ] My change requires a change to the documentation.
[ ] I have updated the tests accordingly (required for a bug fix or a new feature).
[x] I have updated the documentation accordingly.
[x] I have reformatted the code and checked that formatting passes (required).
[x] I have have fixed all issues found by flake8 (required).
[x] I have ensured that all pytest tests pass (required).
[x] I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
[ ] I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line