Run Trax on your own video card w/o colab

Took me a couple of days to hunt down all the pieces and versions, but I now have trax running on my linux desktop. (I have a dual-boot with windows for games, linux for everything else.) Here’s my final solution using a Conda virtual environment under Linux Mint/Ubuntu) if you want to use trax for your own projects:

Create a conda environment and install jupyterlab for the multipage notebook server

conda create -n trax_environment jupyterlab

Install cudatoolkit into the trax environment from channel hcc which has the right Cuda version, 11.2.2 (Cudatoolkit ::

conda activate trax_environment
conda install -c hcc cudatoolkit

Install tools to use the GPU with pip in the trax environment so it can sync the right versions together because conda doesn’t have trax yet

pip install jax tensorflow-gpu trax

upgrade jax to the cuda wheel so it all links correctly for running on your local nvidia GPU

pip install --upgrade "jax[cuda]" -f

To launch the notebook server from the environment:


Now open a notebook and start coding!


@Ken_Otwell What machine and config you have for your dual boot machine? Is it laptop or desktop?

I’m running Mint Cinnamon (Ubuntu) Linux. I was on 20.2 when I did this a few days ago, and just upgraded to 20.3, the very latest version. Still works fine. My machine is AMD 3900x on an Asus X570i motherboard, with an RTX 2080ti which I bought when it first came out because of so much gpu memory (11 gb).

From the system report:
System: Kernel: 5.4.0-96-generic x86_64 bits: 64 compiler: gcc v: 9.3.0 Desktop: Cinnamon 5.2.7
wm: muffin dm: GDM3, LightDM Distro: Linux Mint 20.3 Una base: Ubuntu 20.04 focal
Machine: Type: Desktop Mobo: ASUSTeK model: ROG STRIX X570-I GAMING v: Rev X.0x serial:
UEFI: American Megatrends v: 1405 date: 11/19/2019
CPU: Topology: 12-Core model: AMD Ryzen 9 3900X bits: 64 type: MT MCP arch: Zen
L2 cache: 6144 KiB

Graphics: Device-1: NVIDIA TU102 [GeForce RTX 2080 Ti Rev. A] vendor: driver: nvidia
v: 470.86 bus ID: 0a:00.0 chip ID: 10de:1e07
Display: x11 server: X.Org 1.20.13 driver: nvidia
unloaded: fbdev,modesetting,nouveau,vesa resolution: 3440x1440~60Hz
OpenGL: renderer: NVIDIA GeForce RTX 2080 Ti/PCIe/SSE2 v: 4.6.0 NVIDIA 470.86
direct render: Yes

That was a trove of information. Really grateful to you for so detailed information. You are a true engineer…ThanQ…