non tpu inference about lwm

[wilson1yan]

Hi, inference should work for GPU as well. Perhaps try these installation instructions to see if it works?

conda create -n lwm python=3.10
pip install -U "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install -r requirements.txt

Then bash scripts/run_sample_image.sh (with paths filled out). I tested on an A100 GPU, CUDA 12.3