Warning on slow compilation of model

Hi,

While running ColabFold locally, I get warnings like the following every time the model is (re)compiled, which is slow.

2022-02-15 16:22:05.982180: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:55]
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Compiling module jit_apply_fn.99150
********************************

This seems due to that the model is compiled on CPU rather than GPU. So I guess this is because something is not properly configured. But I didn't find any related configuration that I can change or set within ColabFold.
Is there any idea?

(I have tried setting the env var as the warning suggests, but the corresponding output is a bunch of files describing JAX behaviors, which I have no idea with.)

Thanks!

The issue names no repository files, tests, or entry points. Start by reproducing the local ColabFold run and inspecting the TensorFlow/XLA compilation device configuration; done means determining whether compilation is incorrectly occurring on the CPU and identifying the relevant configuration or cause.