Run on GPU

[Irnamosa]

Hi

Is there any way I could run this repo on GPU instead of CPU for some speed-up?

[Eric Christiansen]

I suspect you could get it working on GPU by just wrapping the relevant part in a tf.device context:
https://www.tensorflow.org/guide/gpu

Otherwise, you might want to find a U-Net variant that's being actively maintained / written using TF 2.0.

[Irnamosa]

Thanks for the help, I managed to get it working!