[Fix] fix microbench for tf2.16 (#2689)

intel · May 8, 2024 · 11c0904 · 11c0904
1 parent 5c043d8
commit 11c0904
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/test/benchmark/test_ApplyAdagradV2.py b/test/benchmark/test_ApplyAdagradV2.py
@@ -17,7 +17,7 @@
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework import dtypes
-from tensorflow.python.training import gen_training_ops
+from tensorflow.python.ops import gen_training_ops
 from utils import multi_run, add_profiling, flush_cache
 from utils import tailed_no_tailed_size
 try:

diff --git a/test/benchmark/test_ApplyGradientDescent.py b/test/benchmark/test_ApplyGradientDescent.py
@@ -17,7 +17,7 @@
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework import dtypes
-from tensorflow.python.training import gen_training_ops
+from tensorflow.python.ops import gen_training_ops
 from utils import multi_run, add_profiling, flush_cache
 from tensorflow.python.framework import constant_op
 from utils import tailed_no_tailed_size