Unable to use multi GPU training with TensorFlow 2 - multi-gpu

Question

Setup: Win 10, 2x Geforce RTX 2080 Ti, Tensorflow 2.9.1 (also tested older versions), Geforce Driver 512.95
I tried multiple tutorials for multi GPU training with Tensorflow 2 and was never able to utelize more then one GPU.
Here is my code.
def mnist_dataset(batch_size):
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train / np.float32(255)
y_train = y_train.astype(np.int64)
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
return train_dataset
def build_and_compile_cnn_model():
model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(512, 3, activation='relu'),
tf.keras.layers.Conv2D(512, 3, activation='relu'),
tf.keras.layers.Conv2D(512, 3, activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10)
])
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.Adam(),
metrics=['accuracy'])
model.summary()
return model
strategy = tf.distribute.MirroredStrategy()
#strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"], cross_device_ops=tf.distribute.ReductionToOneDevice())
print(tf.__version__)
print(tf.config.list_physical_devices())
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
with strategy.scope():
multi_worker_model = build_and_compile_cnn_model()
batch_size = 64
multi_worker_dataset = mnist_dataset(batch_size)
multi_worker_model = build_and_compile_cnn_model()
multi_worker_model.fit(multi_worker_dataset, epochs=10, steps_per_epoch=250)
Output from tf.config.list_physical_devices() and strategy.num_replicas_in_sync.
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')]
Number of devices: 2
But only one GPU gets used.
GPU usage and temperature multi GPU test
To test if both GPUs are working and are accessable by Tensorflow I tried this code.
with tf.device('/gpu:0'):
batch_size = 64
single_worker_dataset = mnist_dataset(batch_size)
single_worker_model = build_and_compile_cnn_model()
single_worker_model.fit(single_worker_dataset, epochs=5, steps_per_epoch=250)
with tf.device('/gpu:1'):
batch_size = 64
single_worker_dataset = mnist_dataset(batch_size)
single_worker_model = build_and_compile_cnn_model()
single_worker_model.fit(single_worker_dataset, epochs=5, steps_per_epoch=250)
Runs as expected. First GPU:0 gets used and then GPU:1. No Error. Everything is fine.
GPU usage and temperature single GPU test
In Tensorflow 1 I was able to use multiple GPUs with keras.utils multi_gpu_model and the same setup.
Any idea what the problem could be?