JIT compilation failed error when running the pix2pix template on GPU

I receive this error when I run the pix2pix template in a jupyter notebook on tensorflow GPU:

This is the fit function, where the error appears:

def fit(train_ds, test_ds, steps):
 example_input, example_target = next(iter(test_ds.take(1)))
 start = time.time()

 for step, (input_image, target) in train_ds.repeat().take(steps).enumerate():
  if (step) % 1000 == 0:

  if step != 0:
    print(f'Time taken for 1000 steps: {time.time()-start:.2f} sec\n')

  start = time.time()

  generate_images(generator, example_input, example_target)
  print(f"Step: {step//1000}k")

train_step(input_image, target, step)

# Training step
if (step+1) % 10 == 0:
  print('.', end='', flush=True)

# Save (checkpoint) the model every 5k steps
if (step + 1) % 5000 == 0:

I have installed cudatoolkit=11.2 cudnn=8.1.0 and tensorflow v2.10 on Windows 10.

When running only with CPU only no error appears though.


Welcome to the Tensorflow Forum!

Thanks for taking the time to report this issue. Since it is a bug, please create an issue on github for the fix.

Thank you very much @chunduriv
I just created an issue on github here