diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index b7bdf3e..815d530 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -2,8 +2,13 @@ FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
 
 WORKDIR /workspace
 
+ARG USERNAME=vscode
+ARG UID=1000
+ARG GID=${UID}
+
 COPY library-scripts/*.sh library-scripts/*.env /tmp/library-scripts/
-RUN apt-get update && export DEBIAN_FRONTEND=noninteractive UID=1000 GID=1000 \
+RUN apt-get update \
+    && export DEBIAN_FRONTEND=noninteractive \
     && /bin/bash /tmp/library-scripts/common-debian.sh \
     && apt-get install -y python3 python3-pip && pip install --upgrade pip --no-input \
     && apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/* /tmp/library-scripts
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index f196787..d4fd3f6 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -4,8 +4,20 @@
   "service": "dev",
   "remoteUser": "vscode",
   "workspaceFolder": "/workspace",
+  "postAttachCommand": "poetry install --with all",
+  "extensions": [
+    "ms-vscode-remote.remote-containers",
+    "ms-azuretools.vscode-docker",
+    "editorconfig.editorconfig",
+    "njpwerner.autodocstring",
+    "ms-python.python",
+    "ms-toolsai.jupyter"
+  ],
+  "runArgs": [
+    "--gpus",
+    "all"
+  ],
   "forwardPorts": [
     8080
-  ],
-  "postAttachCommand": "poetry install --with all"
+  ]
 }
diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml
index 76066ec..e5b0a50 100644
--- a/.devcontainer/docker-compose.yml
+++ b/.devcontainer/docker-compose.yml
@@ -1,18 +1,6 @@
 version: "3"
 
 services:
-
-  # development container
-  dev:
-    container_name: dev
-    build:
-      context: .
-      dockerfile: Dockerfile
-    volumes:
-      - ..:/workspace
-    stdin_open: true
-    # network_mode: service:wandb
-
   # wandb dashboard
   wandb:
     hostname: wandb-local
@@ -23,3 +11,20 @@ services:
     ports:
       - 8080:8080
     restart: unless-stopped
+
+  # development container
+  dev:
+    container_name: dev
+    build:
+      context: .
+      dockerfile: Dockerfile
+    volumes:
+      - ..:/workspace
+    stdin_open: true
+    network_mode: service:wandb
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities:
+                - gpu
diff --git a/.devcontainer/wandb-local/.gitkeep b/.devcontainer/wandb-local/.gitkeep
new file mode 100644
index 0000000..e69de29