Please see my post below for your requested versions.

This same versioning and code almost worked for me - but needs to be changed to: with tf.device("/gpu:0"): model.compile(loss="categorical_crossentropy", optimizer='adma', metrics=['accuracy']) history = model.fit(train_ds,epochs=epochs,validation_data=val_ds, use_multiprocessing=True) Rational: We want to compile the model to use GPU, but keep CPU for image preprocessing. By moving model.fit outside the GPU, it allows both to be used when specified by the user.