message: Add a bunch of CPU blocks to dataset_test
diff: @@ -68,6 +68,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  def testAsFunctionWithMap(self):
  if not context.executing_eagerly():
  self.skipTest("Only works executing eagerly")
  + with ops.device("CPU"):
  original_dataset = dataset_ops.Dataset.range(5).map(lambda x: x * 2)
  fn = original_dataset._trace_variant_creation()
  variant = fn()
  @@ -79,6 +80,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  def testAsFunctionWithMapInFlatMap(self):
  if not context.executing_eagerly():
  self.skipTest("Only works executing eagerly")
  + with ops.device("CPU"):
  original_dataset = dataset_ops.Dataset.range(5).flat_map(
  lambda x: dataset_ops.Dataset.range(5).map(lambda x: x * 2))
  fn = original_dataset._trace_variant_creation()
  @@ -357,6 +359,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  accumulator += value
  return accumulator
  + with ops.device("CPU"):
  first_dataset = dataset_ops.Dataset.range(10)
  self.assertEqual(45, self.evaluate(_uses_dataset(first_dataset)))
  second_dataset = dataset_ops.Dataset.range(11)
  