Please use Tensor Flow TensorFlow Version 2.7.0
As mentioned by Frameworks Engineer & chaitanya_ambaruk, please use the with tf.device('/CPU:0'):
PS: Need to place this piece of code in the correct line of code. For the above example please place it before you start the augmentation process
with tf.device('/CPU:0'):
data_augmentation = keras.Sequential(
[
layers.experimental.preprocessing.RandomFlip("horizontal_and_vertical", input_shape = (img_height, img_width,3)),
layers.experimental.preprocessing.RandomRotation(0.2),
layers.experimental.preprocessing.RandomZoom(0.2)
]
)
plt.figure(figsize=(10,10))
for images, _ in train_ds.take(5):
for i in range(9):
augmented_images = data_augmentation(images)
ax = plt.subplot(3,3,i+1)
plt.imshow(augmented_images[0].numpy().astype('uint8'))
plt.axis("off");