Set TF_FORCE_GPU_ALLOW_GROWTH=true by default

[samos123]

This is needed to be able to run Fuji v2 70B on GPU without GPU memory OOMs.

@kelvin-zou can likely confirm whether this should be the default or not.

[kelvin-zou]

Thanks!
Can we move it to somewhere since it is GPU specific? This launch cmd is shared across GPU and TPU.

[samos123]

Hmm but we also set a lot of TPU environment variables in launch.py without any if statements. I don't think there is a better place since it needs to happen before jax is started?

Would you prefer this?

if instance_type.startswith("gpu"):
    # Prevent GPU OOM issues due to TF taking up all the GPU memory.
    # Reference: https://stackoverflow.com/a/54927279
    os.environ.setdefault("TF_FORCE_GPU_ALLOW_GROWTH", "true")