How to accelerate the encoding time when using jax-unirep?

[wushufang]

Hi, Is there any method to accelerate the encoding time? I do not need to fine tune the model. It is too slow now.

[Binyun-Z]

I have the same problem,when I want to get embedding for some proteins .it shows "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)",however, there is TeslaV100 GPU in my PC