commit 52a7c2d038a9440b35a4d7f64f2846f620592a5d
author mdfaijul <md.faijul.amin@intel.com> Mon Oct 25 16:36:53 2021 -0700
committer mdfaijul <md.faijul.amin@intel.com> Tue Oct 26 08:32:53 2021 -0700
tree 90db8072d2e4156b127c3daae943fbcf98ee892b
parent 3ca3f16c9b10aa0ebd6703614345e2ef8210eb0d

Add CPU device.

tensorflow/python/grappler/remapper_test.py

diff --git a/tensorflow/python/grappler/remapper_test.py b/tensorflow/python/grappler/remapper_test.py
index 5cc9ec2..6aff449 100644
--- a/tensorflow/python/grappler/remapper_test.py
+++ b/tensorflow/python/grappler/remapper_test.py
@@ -96,16 +96,17 @@
 
         # Create MatMul + BiasAdd + Gelu graph
         ops.reset_default_graph()
-        x = _input([m, k])
-        w = _weight([k, n])
-        b = _bias([n])
-        if precision == 'bfloat16':
-          x = math_ops.cast(x, dtypes.bfloat16)
-          w = math_ops.cast(w, dtypes.bfloat16)
-          b = math_ops.cast(b, dtypes.bfloat16)
-        y = math_ops.matmul(x, w)
-        z = nn.bias_add(y, b)
-        out = nn.gelu(z, approximate=approximate)
+        with ops.device('/device:CPU:0'):
+          x = _input([m, k])
+          w = _weight([k, n])
+          b = _bias([n])
+          if precision == 'bfloat16':
+            x = math_ops.cast(x, dtypes.bfloat16)
+            w = math_ops.cast(w, dtypes.bfloat16)
+            b = math_ops.cast(b, dtypes.bfloat16)
+          y = math_ops.matmul(x, w)
+          z = nn.bias_add(y, b)
+          out = nn.gelu(z, approximate=approximate)
 
         # Compute reference value.
         config = _get_config(remapping_on=False)