Seems that Tensorflow has been pretty much abandoned gpu support on Windows and Jax is used internally at Google, is there any plans to just switch to PyTorch? I am also unable to get my gpu to load the mosaic model even with the --is_mosaic=True argument.