Installation instruction does not mention which CUDA or NVIDIA drivers are required

[SupermicroML]

When we follow the instructions below

_pip3 install -e '.[pytorch_cpu]'
pip3 install -e '.[jax_gpu]' -f 'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html'
pip3 install -e '.[full]'
TL;DR to install the PyTorch version for GPU run:

pip3 install -e '.[jax_cpu]'
pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'
pip3 install -e '.[full]'_

They automatically install torch for CUDA 11.8. However the host system may have different CUDA version and NVIDIA driver at the time of installation. Perhaps it is good to add a seciton "Prerequsites" and add CUDA version details inside. This will prevent mismatch between package versions. The same is also valid for the Docker build steps.

For example, if the developers have CUDA > 12.0, can they still install the package or they need to downgrade CUDA to 11.8?

Best Regards,
Nikola

[priyakasimbeg]

Hi Nikola,
Thanks for the suggestion! We will document the driver and CUDA version.

The same is also valid for Docker build steps.

The CUDA version (11.8) is built into the Docker image and your global system level CUDA installation does not affect anything in the docker build process.

For example, if the developers have CUDA > 12.0, can they still install the package or they need to downgrade CUDA to 11.8?

If developers have a different version of CUDA we recommend running the container build from the image in docker/Dockerfile.

Does this address your concerns?

[priyakasimbeg]

Updated documentation in #570