diff --git a/apex/amp/frontend.py b/apex/amp/frontend.py index cbaf139dc..b97e5e796 100644 --- a/apex/amp/frontend.py +++ b/apex/amp/frontend.py @@ -1,4 +1,5 @@ import torch +import os from ._initialize import _initialize from ._amp_state import _amp_state, warn_or_err, maybe_print from collections import OrderedDict @@ -422,6 +423,11 @@ def initialize( for k, v in _amp_state.opt_properties.options.items(): maybe_print("{:26} : {}".format(k, v), True) + + # Set flag to tell others apex.amp is enabled + os.environ["APEX_AMP_ENABLED"] = "1" + + return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)