diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 1484088..25bbb77 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -7,6 +7,8 @@ on:
 jobs:
   deploy:
     runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/catthehacker/ubuntu:runner-latest
     timeout-minutes: 30
     steps:
       - uses: actions/checkout@v2
diff --git a/.github/workflows/tests_linters.yml b/.github/workflows/tests_linters.yml
index b88b89e..3b6eee9 100644
--- a/.github/workflows/tests_linters.yml
+++ b/.github/workflows/tests_linters.yml
@@ -4,26 +4,22 @@ on: [ push, pull_request ]
 
 jobs:
   tests-and-linters:
-    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
-    runs-on: "${{ matrix.os }}"
-
-    strategy:
-      matrix:
-        python-version: ["3.9"]
-        os: [self-hosted]
+    name: "Python 3.9 on GitHub Hosted runner"
+    runs-on: ubuntu-latest
+    container:
+        image: python:3.9
 
     steps:
       - name: Install dependencies for viewer test
-        run: sudo apt-get update && sudo apt-get install -y xvfb
+        run: apt-get update && apt-get install -y xvfb
       - name: Checkout flashbax
         uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-            python-version: "${{ matrix.python-version }}"
       - name: Install python dependencies 🔧
         run: pip install .[dev]
       - name: List python packages 📦
         run: pip list
+      - name: Update git permissions
+        run: git config --global --add safe.directory /__w/flashbax/flashbax
       - name: Run linters 🖌️
         run: pre-commit run --all-files --verbose
       - name: Run tests 🧪
diff --git a/flashbax/buffers/prioritised_trajectory_buffer.py b/flashbax/buffers/prioritised_trajectory_buffer.py
index 88893a3..0dae6fb 100644
--- a/flashbax/buffers/prioritised_trajectory_buffer.py
+++ b/flashbax/buffers/prioritised_trajectory_buffer.py
@@ -798,6 +798,7 @@ def make_prioritised_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         prioritised_init,
         add_batch_size=add_batch_size,
diff --git a/flashbax/buffers/trajectory_buffer.py b/flashbax/buffers/trajectory_buffer.py
index ba005cd..e280420 100644
--- a/flashbax/buffers/trajectory_buffer.py
+++ b/flashbax/buffers/trajectory_buffer.py
@@ -585,6 +585,7 @@ def make_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         init,
         add_batch_size=add_batch_size,