Running chroma from docker

generatebio / chroma

A generative model for programmable protein design

Apache License 2.0

627 stars 74 forks source link

Closed e-osipov closed 7 months ago

e-osipov commented 7 months ago

Hi everyone,

I wonder how to properly create docker image and run container in order to utilize GPU? For now, chroma is running using CPU only.

What I did:

I created an image from Dockerfile using docker build -t chroma . command
Container was started using docker run -i -t --runtime=nvidia --gpus all chroma
nvidia-smi within container will show that I have GPU available However, any example will run on CPU, using all available cores.

wujiewang commented 7 months ago

Ok, I think torch is somehow not installed properly with cuda support, this is weird....

I will update the Dockerfile to explicitly specify torch version with cuda=11.3.

Sorry about that!

wujiewang commented 7 months ago

Oh, I guess cuda=11.3.1 is too old for torch=2.1.1, will just do

pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

wujiewang commented 7 months ago

ok, resolved with d5f5787!! thanks for reporting this issue!