Installing Frameworks (PyTorch, TensorFlow, Jax)
This page describes installing popular frameworks like TensorFlow, PyTorch & JAX, etc. on your Oscar account.
In this example, we will install PyTorch (refer to sub-pages for TensorFlow and Jax).
Step 1: Request an interactive session on a GPU node with Ampere architecture GPUs
interact -q gpu -g 1 -f ampere -m 20g -n 4
Here, -f = feature. We only need to build on Ampere once.
Step 2: Once your session has started on a compute node, run nvidia-smi
to verify the GPU and then load the appropriate modules
Step 3: Create and activate the virtual environment, unload the pre-loaded modules then load cudnn and cuda dependencies
Step 4: Create a new vittual environment
Step 5: Install the required packages
The aforementioned will install the latest version of PyTorch with cuda11 compatibility, for older versions you can specify the version by:
Step 6: Test that PyTorch is able to detect GPUs
If the above functions return True
and GPU model
, then it's working correctly. You are all set, now you can install other necessary packages.
Last updated
Was this helpful?