diff --git a/dev/.devcontainer/Dockerfile b/dev/.devcontainer/Dockerfile
index 11fb57e0b..6f48300f7 100644
--- a/dev/.devcontainer/Dockerfile
+++ b/dev/.devcontainer/Dockerfile
@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See https://go.microsoft.com/fwlink/?linkid=2090316 for license information.
 #-------------------------------------------------------------------------------------------------------------
 
-FROM python:3.7
+FROM python:3.12
 
 # Avoid warnings by switching to noninteractive
 ENV DEBIAN_FRONTEND=noninteractive
@@ -52,7 +52,7 @@ RUN apt-get update \
 RUN pip install numpy jaxlib tensorflow tensorflow-datasets matplotlib msgpack \
     jupyter pytest pytest-xdist \
     twine \
-    sphinx sphinx_rtd_theme ipykernel nbsphinx recommonmark sklearn
+    sphinx sphinx_rtd_theme ipykernel nbsphinx recommonmark scikit-learn
 
 # Switch back to dialog for any ad-hoc use of apt-get
 ENV DEBIAN_FRONTEND=dialog
diff --git a/dev/.devcontainer/devcontainer.json b/dev/.devcontainer/devcontainer.json
index 32b87efa7..2b6cde283 100644
--- a/dev/.devcontainer/devcontainer.json
+++ b/dev/.devcontainer/devcontainer.json
@@ -22,7 +22,7 @@
 	// Uncomment the next line to run commands after the container is created.
 	// "postCreateCommand": "pip install -r requirements.txt",
 	"runArgs": ["-v", "${env:HOME}${env:USERPROFILE}/.ssh:/root/.ssh-localhost:ro"],
-	"postCreateCommand": "sudo cp -r /root/.ssh-localhost ~/.ssh && sudo chown -R $(id -u):$(id -g) ~/.ssh && chmod 700 ~/.ssh && chmod 600 ~/.ssh/* && pip install -e ./flax",
+	"postCreateCommand": "sudo cp -r /root/.ssh-localhost ~/.ssh && sudo chown -R $(id -u):$(id -g) ~/.ssh && chmod 700 ~/.ssh && chmod 600 ~/.ssh/* && pip install -e .",
 
 	// Uncomment the next line to have VS Code connect as an existing non-root user in the container.
 	// On Linux, by default, the container user's UID/GID will be updated to match your local user. See
diff --git a/dev/README.md b/dev/README.md
index 5d8f081bb..56f3484b5 100644
--- a/dev/README.md
+++ b/dev/README.md
@@ -13,4 +13,7 @@ the environments used by contributors and maintainers.
 5. Re-open the folder workspace using the remote containers extension.
    VSCode should recommend this action in a popup. Alternatively,
    use the green button in the bottom left container to control the
-   remote extension.
\ No newline at end of file
+   remote extension.
+
+## Troubleshoot:
+If you have the following error `~/.docker/buildx/current: permission denied`, try running `sudo chown -R $(whoami) ~/.docker`
\ No newline at end of file