Add NVIDIA apex support and gradient checkpointing to reduce memory footprint

[seovchinnikov]

I've added NVIDIA apex support and checkpointing (https://pytorch.org/docs/stable/checkpoint.html) mechanism to reduce memory footprint.

You can run it with --checkpointing --opt_level "O2" and increased input crop size (I was able to run CycleGAN with up to 896 on my 2080 RTX). Checkpointing is only used for CycleGAN for now (can be improved further).
Please note that it was tested on pytorch 1.7 nightly build, and behavior of apex is unstable on old versions.

[junyanz]

Great feature! I am wondering if you can get the same results with and without apex and gradient checkpointing.

[seovchinnikov]

I think we should run base tests and check it against the baselines

[vict0rsch]

Note that amp is part of pytorch as of 1.6 => https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples