Update Intel_TensorFlow_AMX_BF16_Inference.py

Author: YuningQiu

AI-and-Analytics/Features-and-Functionality/IntelTensorFlow_AMX_BF16_Inference/Intel_TensorFlow_AMX_BF16_Inference.py

@@ -80,7 +80,7 @@
 
 
 # Reload the model as the bf16 model with AVX512 to compare inference time
-os.environ["ONEDNN_MAX_CPU_ISA"] = "AVX512_BF16"
+os.environ["ONEDNN_MAX_CPU_ISA"] = "AVX512_CORE_BF16"
 tf.config.optimizer.set_experimental_options({'auto_mixed_precision_onednn_bfloat16':True})
 bf16_model_noAmx = tf.keras.models.load_model('models/my_saved_model_fp32')
 
@@ -93,7 +93,7 @@
 
 
 # Reload the model as the bf16 model with AMX to compare inference time
-os.environ["ONEDNN_MAX_CPU_ISA"] = "AMX_BF16"
+os.environ["ONEDNN_MAX_CPU_ISA"] = "AVX512_CORE_AMX"
 tf.config.optimizer.set_experimental_options({'auto_mixed_precision_onednn_bfloat16':True})
 bf16_model_withAmx = tf.keras.models.load_model('models/my_saved_model_fp32')