redefine enable_auto_mix_precision api

Author: XiaobingSuper

intel_pytorch_extension_py/__init__.py

@@ -82,16 +82,19 @@ def enable_auto_optimization(mixed_dtype = None, train = False):
     """
     if mixed_dtype != None:
         core.enable_auto_dnnl()
+    enable_auto_mix_precision(mixed_dtype, train)
+
+def enable_auto_mix_precision(mixed_dtype = torch.bfloat16, train = False):
     running_mode = 'training' if train else 'inference'
-    enable_auto_mix_precision(AmpConf(mixed_dtype), running_mode).__enter__()
+    auto_mix_precision(AmpConf(mixed_dtype), running_mode).__enter__()
 
 def get_auto_optimization():
     return get_auto_mix_precision
 
 def get_train():
     return core.get_train()
 
-class enable_auto_mix_precision(_DecoratorContextManager):
+class auto_mix_precision(_DecoratorContextManager):
     def __init__(self, conf, running_mode = 'inference'):
         self.pre_mixed_dtype = get_auto_mix_precision()
         self.pre_running_mode = get_train()

tests/cpu/common_ipex_conf.py

@@ -4,21 +4,21 @@
 class AutoMixPrecision(object):
     def __init__(self, enable_or_not = False, train = False):
         self.old_value = ipex.get_auto_mix_precision()
-        self.pre_running_mode = 'training' if ipex.get_train() else 'inference'
+        self.train_old_value = ipex.get_train()
         self.enable_or_not = enable_or_not
-        self.running_mode = 'training' if train else 'inference'
+        self.train = train
 
     def __enter__(self):
         if self.enable_or_not:
-            ipex.enable_auto_mix_precision(ipex.AmpConf(torch.bfloat16), self.running_mode).__enter__()
+            ipex.enable_auto_mix_precision(mixed_dtype=torch.bfloat16, train=self.train)
         else:
-            ipex.enable_auto_mix_precision(ipex.AmpConf(None)).__enter__()
+            ipex.enable_auto_mix_precision(mixed_dtype=None)
 
     def __exit__(self, *args, **kwargs):
         if self.old_value:
-            ipex.enable_auto_mix_precision(ipex.AmpConf(torch.bfloat16), self.pre_running_mode).__enter__()
+            ipex.enable_auto_mix_precision(mixed_dtype=torch.bfloat16, train=self.train_old_value)
         else:
-            ipex.enable_auto_mix_precision(ipex.AmpConf(None)).__enter__()
+            ipex.enable_auto_mix_precision(mixed_dtype=None)
 
 class AutoDNNL(object):
     def __init__(self, enable_or_not = False):