flax

b5ca503e - Hide GPUs from TensorFlow and use default v2.

Hide GPUs from TensorFlow and use default v2. This change hides GPUs from TensorFlow to avoid it from grabbing GPU memory (see issue #445). Also bumps to TensorFlow 2 has been the default for a while.