cleanup

syrkis · Jun 9, 2024 · ca9ab84 · ca9ab84
1 parent 9eb9550
commit ca9ab84
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -19,8 +19,9 @@ COPY requirements.txt .
 
 RUN python3.11 -m pip install -r requirements.txt
 
-RUN python3.11 -m pip install \
-    "jax[cuda11]"
+RUN python3.11 -m pip install --upgrade \
+    "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
 
 ENV PYGLFW_PREVIEW=1
 

diff --git a/Makefile b/Makefile
diff --git a/main.ua b/main.ua
diff --git a/src/plots.py b/src/plots.py
@@ -94,6 +94,8 @@ def curve_plot(
     fig, ax = init_curve_plot(info, conf)
     for i, curve in enumerate(curves.T):  # transpose bcs of jax.lax.scan
         ax.plot(curve, c=ink, lw=2, ls="--" if i > 0 else "-")
+    # make x-axis log
+    ax.set_xscale("log")
     # add info key, val pairs to right side of plot (centered verticall on right axis) (relative to len(info))
     ignore = ["in_d", "out_d", "block"]
     conf = {k: v for k, v in conf.items() if k not in ignore}