[training] add an offload utility that can be used as a context manager.

[sayakpaul]

What does this PR do?

Useful for handling offloading in a simple and easy manner.

@DN6 I am including changes to examples/**/*.py to also trigger the fast PR tests on GPU. Hope, that's okay?

Cc: @a-r-r-o-w for visibility.