diff --git a/.github/workflows/python-checks.yml b/.github/workflows/python-checks.yml
index 6af86284..d52709c5 100644
--- a/.github/workflows/python-checks.yml
+++ b/.github/workflows/python-checks.yml
@@ -62,7 +62,7 @@ jobs:
 
       - name: Install Lightning dependency
         run: |
-          python -m pip install -e "s3torchconnector[lightning]"
+          python -m pip install -e "s3torchconnector[lightning-tests]"
       - name: s3torchconnector lightning unit tests
         run: pytest s3torchconnector/tst/unit/lightning --hypothesis-profile ci --hypothesis-show-statistics -c ./
 
diff --git a/.github/workflows/python-integration.yml b/.github/workflows/python-integration.yml
index 8af5e730..6ce90b4d 100644
--- a/.github/workflows/python-integration.yml
+++ b/.github/workflows/python-integration.yml
@@ -93,7 +93,7 @@ jobs:
 
       - name: Install Lightning dependency
         run: |
-          python -m pip install -e "s3torchconnector[lightning]"
+          python -m pip install -e "s3torchconnector[lightning-tests]"
       - name: s3torchconnector lightning e2e tests
         run: |
           CI_REGION=${{ matrix.test-run.region }} \
diff --git a/s3torchconnector/pyproject.toml b/s3torchconnector/pyproject.toml
index 8e830080..4beb233f 100644
--- a/s3torchconnector/pyproject.toml
+++ b/s3torchconnector/pyproject.toml
@@ -46,7 +46,11 @@ e2e = [
 ]
 
 lightning = [
-    "lightning >= 2.0",
+    "lightning >= 2.0"
+]
+
+lightning-tests = [
+    "s3torchconnector[lightning]",
     "s3fs"
 ]
 
diff --git a/s3torchconnector/tst/e2e/models/lightning_transformer.py b/s3torchconnector/tst/e2e/models/lightning_transformer.py
index f75b3329..a197756e 100644
--- a/s3torchconnector/tst/e2e/models/lightning_transformer.py
+++ b/s3torchconnector/tst/e2e/models/lightning_transformer.py
@@ -1,10 +1,14 @@
+"""
+Based on: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+Source: https://github.com/Lightning-AI/pytorch-lightning/blob/master/docs/source-pytorch/common/lightning_module.rst
+License: https://github.com/Lightning-AI/pytorch-lightning/blob/master/LICENSE
+"""
+
 import lightning as L
 import torch
 
 from lightning.pytorch.demos import Transformer
 
-# Based on https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
-
 
 class LightningTransformer(L.LightningModule):
     def __init__(self, vocab_size):
diff --git a/s3torchconnector/tst/e2e/models/net.py b/s3torchconnector/tst/e2e/models/net.py
index 89608f0b..9de0d700 100644
--- a/s3torchconnector/tst/e2e/models/net.py
+++ b/s3torchconnector/tst/e2e/models/net.py
@@ -1,9 +1,13 @@
+"""
+Based on: https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html
+Source: https://github.com/pytorch/tutorials/blob/main/beginner_source/blitz/neural_networks_tutorial.py
+License: https://github.com/pytorch/tutorials/blob/main/LICENSE
+"""
+
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-# Based on https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html
-
 
 class Net(nn.Module):
     def __init__(self):
diff --git a/s3torchconnectorclient/pyproject.toml b/s3torchconnectorclient/pyproject.toml
index 14b033cd..651e07fc 100644
--- a/s3torchconnectorclient/pyproject.toml
+++ b/s3torchconnectorclient/pyproject.toml
@@ -53,7 +53,7 @@ test-command = [
     "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} pytest {package}/../s3torchconnector/tst/e2e --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py",
     "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' pytest {package}/python/tst/integration",
     "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' pytest {package}/../s3torchconnector/tst/e2e  --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py",
-    "python -m pip install -e '{package}/../s3torchconnector[lightning]'",
+    "python -m pip install -e '{package}/../s3torchconnector[lightning-tests]'",
     "pytest {package}/../s3torchconnector/tst/unit/lightning",
     "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} pytest {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py",
     "AWS_DEFAULT_REGION=${S3_EXPRESS_REGION} CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' pytest {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py",