diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 102a8e0c5..852e7dccf 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -16,7 +16,8 @@ jobs:
         python -m pip install --upgrade pip
         pip install pytest
         pip install .[jax_cpu]
+    - name: Setup tmate session
+      uses: mxschmitt/action-tmate@v3
     - name: Run pytest
       run: |
-
         pytest -vx tests/
diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml
index 556fd0f99..b880bde10 100644
--- a/.github/workflows/linting.yml
+++ b/.github/workflows/linting.yml
@@ -20,7 +20,7 @@ jobs:
         pylint algorithmic_efficiency
         pylint baselines
         pylint submission_runner.py
-        pylint submission_runner_test.py
+        pylint tests
 
   isort:
     runs-on: ubuntu-latest
diff --git a/algorithmic_efficiency/workloads/ogbg/ogbg_jax/workload.py b/algorithmic_efficiency/workloads/ogbg/ogbg_jax/workload.py
index 00c8518b5..1ffb47326 100644
--- a/algorithmic_efficiency/workloads/ogbg/ogbg_jax/workload.py
+++ b/algorithmic_efficiency/workloads/ogbg/ogbg_jax/workload.py
@@ -9,10 +9,10 @@
 
 from algorithmic_efficiency import random_utils as prng
 from algorithmic_efficiency import spec
-from algorithmic_efficiency.workloads.ogbg.workload import BaseOgbgWorkload
 from algorithmic_efficiency.workloads.ogbg.ogbg_jax import input_pipeline
 from algorithmic_efficiency.workloads.ogbg.ogbg_jax import metrics
 from algorithmic_efficiency.workloads.ogbg.ogbg_jax import models
+from algorithmic_efficiency.workloads.ogbg.workload import BaseOgbgWorkload
 
 
 class OgbgWorkload(BaseOgbgWorkload):