Add setting of env flag when apex is turned on

ROCm · Sep 5, 2023 · 9443e1a · 9443e1a
1 parent 8fc9b21
commit 9443e1a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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)