Fix colab issues

[fm1320]

What does this PR do?

fixes #294 and classification optimization notebook

Fixes #<issue_number>

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB