O1 opt_level patches pytorch calls outside wrapped modules.

If you initialize some (but not all) modules with AMP opt level O1, all pytorch calls are now patched, even those outside wrapped modules. I'm not sure if this is intentional or a bug, but the behavior was definitely unexpected.

(to be clear, I know that O1 patches torch functions, but the fact that it patches torch functions outside modules passed to amp.initialize() was the surprise)

Example

import torch as th
from apex import amp

f = th.nn.Conv2d(3, 3, 3, 1, 1).cuda()
g = th.nn.Conv2d(3, 3, 3, 1, 1).cuda()
x = th.zeros(1, 3, 32, 32).cuda()
print(f(x).dtype, g(x).dtype)

f, = amp.initialize([f], opt_level="O1")
print(f(x).dtype, g(x).dtype)

prints torch.float16 torch.float16 even though g was not wrapped.

Start by reproducing the provided example with AMP opt level O1, comparing the dtypes of calls through f and the unwrapped g. Trace the O1 initialization and patching behavior to determine whether calls outside modules passed to amp.initialize() should be affected. Done means the behavior is clarified and covered by an appropriate regression test.