gistfile1.txt

$ docker run --gpus all -i -t jax:latest /bin/bash
(base) [docker@e697aef58065 ~]$ ls
anaconda  cuda-repo-rhel8-10-2-local-10.2.89-440.33.01-1.0-1.x86_64.rpm
(base) [docker@e697aef58065 ~]$ which python
~/anaconda/bin/python
(base) [docker@e697aef58065 ~]$ conda activate mouse-hmm
(mouse-hmm) [docker@e697aef58065 ~]$ python
Python 3.7.6 | packaged by conda-forge | (default, Jun  1 2020, 18:57:50) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax.numpy as np
>>> np.ones(3)
2020-06-26 01:10:42.256378: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:314] failed call to cuInit: CUDA_ERROR_SYSTEM_DRIVER_MISMATCH: system has unsupported display driver / cuda driver combination
2020-06-26 01:10:42.256698: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_diagnostics.cc:313] kernel version 440.36.0 does not match DSO version 440.33.1 -- cannot find working devices in this configuration
/home/docker/anaconda/envs/mouse-hmm/lib/python3.7/site-packages/jax/lib/xla_bridge.py:125: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
DeviceArray([1., 1., 1.], dtype=float32)
>>> 
(mouse-hmm) [docker@e697aef58065 ~]$ nvidia-smi 
Fri Jun 26 01:11:19 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.36       Driver Version: 440.36       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  Off  | 00000000:04:00.0 Off |                  N/A |
|  0%   24C    P8    12W / 250W |    318MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+
(mouse-hmm) [docker@e697aef58065 ~]$