diff --git a/.conda/meta.yaml b/.conda/meta.yaml
index 3a723a9..c227313 100644
--- a/.conda/meta.yaml
+++ b/.conda/meta.yaml
@@ -1,30 +1,26 @@
-# https://docs.conda.io/projects/conda-build/en/latest/resources/define-metadata.html#loading-data-from-other-files
-# https://github.com/conda/conda-build/pull/4480
-# for conda-build > 3.21.9
-# {% set pyproject = load_file_data('../pyproject.toml', from_recipe_dir=True) %}
-# {% set project = pyproject.get('project') %}
-# {% set urls = pyproject.get('project', {}).get('urls') %}
+{% set pyproject = load_file_data('../pyproject.toml', from_recipe_dir=True) %}
+{% set project = pyproject.get('project') %}
+{% set urls = pyproject.get('project', {}).get('urls') %}
+{% set version = environ.get('BUILD_VERSION', '0.2.0.dev0') %}
 package:
-  name: torchscan
-  version: "{{ environ.get('BUILD_VERSION') }}"
+  name: {{ project.get('name') }}
+  version: {{ version }}
 
 source:
-  fn: torchscan-{{ environ.get('BUILD_VERSION') }}.tar.gz
-  url: ../dist/torchscan-{{ environ.get('BUILD_VERSION') }}.tar.gz
+  fn: {{ project.get('name') }}-{{ version }}}.tar.gz
+  url: ../dist/{{ project.get('name') }}-{{ version }}.tar.gz
 
 build:
-  number: 0
   noarch: python
   script: python setup.py install --single-version-externally-managed --record=record.txt
 
 requirements:
   host:
-    - python>=3.6, <4.0
+    - python>=3.8, <4.0
     - setuptools
 
   run:
-    - python>=3.6, <4.0
-    - pytorch >=1.5.0, <2.0.0
+    - pytorch >=2.0.0, <3.0.0
 
 test:
   # Python imports
@@ -33,13 +29,15 @@ test:
     - torchscan.modules
     - torchscan.process
     - torchscan.utils
+  requires:
+    - python
 
 about:
-  home: https://github.com/frgfm/torch-scan
+  home: {{ urls.get('repository') }}
   license: Apache 2.0
-  license_file: LICENSE
-  summary: 'Useful information about your Pytorch module'
+  license_file: {{ project.get('license', {}).get('file') }}
+  summary: {{ project.get('description') }}
   # description: |
   #   {{ data['long_description'] | replace("\n", "\n    ") | replace("#", '\#')}}
-  doc_url: https://frgfm.github.io/torch-scan/
-  dev_url: https://github.com/frgfm/torch-scan
+  doc_url: {{ urls.get('documentation') }}
+  dev_url: {{ urls.get('repository') }}
diff --git a/.flake8 b/.flake8
deleted file mode 100644
index 70ebcb5..0000000
--- a/.flake8
+++ /dev/null
@@ -1,5 +0,0 @@
-[flake8]
-max-line-length = 120
-ignore = E203, E402, E265, F403, W503, W504, E731
-exclude = .git, venv*, build, docs
-per-file-ignores = **/__init__.py:F401
diff --git a/.github/collect_env.py b/.github/collect_env.py
index 961510e..4664c4f 100644
--- a/.github/collect_env.py
+++ b/.github/collect_env.py
@@ -14,9 +14,10 @@
 import locale
 import os
 import re
-import subprocess
+import subprocess  # noqa S404
 import sys
-from collections import namedtuple
+from pathlib import Path
+from typing import NamedTuple
 
 try:
     import torchscan
@@ -36,20 +37,16 @@
 
 
 # System Environment Information
-SystemEnv = namedtuple(
-    "SystemEnv",
-    [
-        "torchscan_version",
-        "torch_version",
-        "os",
-        "python_version",
-        "is_cuda_available",
-        "cuda_runtime_version",
-        "nvidia_driver_version",
-        "nvidia_gpu_models",
-        "cudnn_version",
-    ],
-)
+class SystemEnv(NamedTuple):
+    torchscan_version: str
+    torch_version: str
+    os: str
+    python_version: str
+    is_cuda_available: bool
+    cuda_runtime_version: str
+    nvidia_driver_version: str
+    nvidia_gpu_models: str
+    cudnn_version: str
 
 
 def run(command):
@@ -125,18 +122,18 @@ def get_cudnn_version(run_lambda):
     # find will return 1 if there are permission errors or if not found
     if len(out) == 0 or rc not in (1, 0):
         lib = os.environ.get("CUDNN_LIBRARY")
-        if lib is not None and os.path.isfile(lib):
+        if lib is not None and Path(lib).is_file():
             return os.path.realpath(lib)
         return None
     files = set()
     for fn in out.split("\n"):
         fn = os.path.realpath(fn)  # eliminate symbolic links
-        if os.path.isfile(fn):
+        if Path(fn).is_file():
             files.add(fn)
     if not files:
         return None
     # Alphabetize the result because the order is non-deterministic otherwise
-    files = list(sorted(files))
+    files = sorted(files)
     if len(files) == 1:
         return files[0]
     result = "\n".join(files)
@@ -149,11 +146,11 @@ def get_nvidia_smi():
     if get_platform() == "win32":
         system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
         program_files_root = os.environ.get("PROGRAMFILES", "C:\\Program Files")
-        legacy_path = os.path.join(program_files_root, "NVIDIA Corporation", "NVSMI", smi)
-        new_path = os.path.join(system_root, "System32", smi)
+        legacy_path = Path(program_files_root) / "NVIDIA Corporation" / "NVSMI" / smi
+        new_path = Path(system_root) / "System32" / smi
         smis = [new_path, legacy_path]
         for candidate_smi in smis:
-            if os.path.exists(candidate_smi):
+            if Path(candidate_smi).exists():
                 smi = '"{}"'.format(candidate_smi)
                 break
     return smi
@@ -220,10 +217,7 @@ def get_os(run_lambda):
 def get_env_info():
     run_lambda = run
 
-    if TORCHSCAN_AVAILABLE:
-        torchscan_str = torchscan.__version__
-    else:
-        torchscan_str = "N/A"
+    torchscan_str = torchscan.__version__ if TORCHSCAN_AVAILABLE else "N/A"
 
     if TORCH_AVAILABLE:
         torch_str = torch.__version__
@@ -261,14 +255,14 @@ def get_env_info():
 
 def pretty_str(envinfo):
     def replace_nones(dct, replacement="Could not collect"):
-        for key in dct.keys():
+        for key in dct:
             if dct[key] is not None:
                 continue
             dct[key] = replacement
         return dct
 
     def replace_bools(dct, true="Yes", false="No"):
-        for key in dct.keys():
+        for key in dct:
             if dct[key] is True:
                 dct[key] = true
             elif dct[key] is False:
@@ -292,7 +286,7 @@ def maybe_start_on_next_line(string):
         "nvidia_gpu_models",
         "nvidia_driver_version",
     ]
-    all_cuda_fields = dynamic_cuda_fields + ["cudnn_version"]
+    all_cuda_fields = [*dynamic_cuda_fields, "cudnn_version"]
     all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None for field in dynamic_cuda_fields)
     if TORCH_AVAILABLE and not torch.cuda.is_available() and all_dynamic_cuda_fields_missing:
         for field in all_cuda_fields:
diff --git a/.github/verify_labels.py b/.github/verify_labels.py
index 143d616..58f6c99 100644
--- a/.github/verify_labels.py
+++ b/.github/verify_labels.py
@@ -45,7 +45,9 @@
 
 
 def query_repo(cmd: str, *, accept) -> Any:
-    response = requests.get(f"https://api.github.com/repos/{GH_ORG}/{GH_REPO}/{cmd}", headers=dict(Accept=accept))
+    response = requests.get(
+        f"https://api.github.com/repos/{GH_ORG}/{GH_REPO}/{cmd}", headers={"Accept": accept}, timeout=5
+    )
     return response.json()
 
 
diff --git a/.github/workflows/builds.yml b/.github/workflows/builds.yml
index a1f9744..6b0180f 100644
--- a/.github/workflows/builds.yml
+++ b/.github/workflows/builds.yml
@@ -13,11 +13,10 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python: [3.7, 3.8]
+        python: [3.8, 3.9, '3.10', 3.11]
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
@@ -29,6 +28,48 @@ jobs:
       - name: Install package
         run: |
           python -m pip install --upgrade pip
-          pip install -e . --upgrade
+          pip install -e .
       - name: Import package
         run: python -c "import torchscan; print(torchscan.__version__)"
+
+  pypi:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+          architecture: x64
+      - name: Cache python modules
+        uses: actions/cache@v2
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-python-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}-build
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install setuptools wheel twine --upgrade
+      - run: |
+          python setup.py sdist bdist_wheel
+          twine check dist/*
+
+  conda:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          auto-update-conda: true
+          python-version: "3.9"
+      - name: Install dependencies
+        shell: bash -el {0}
+        run: conda install -y conda-build conda-verify
+      - name: Build conda
+        shell: bash -el {0}
+        run: |
+          python setup.py sdist
+          mkdir conda-dist
+          conda env list
+          conda build .conda/ -c pytorch --output-folder conda-dist
+          ls -l conda-dist/noarch/*tar.bz2
diff --git a/.github/workflows/doc-status.yml b/.github/workflows/doc-status.yml
index ef9a250..d2f2d9d 100644
--- a/.github/workflows/doc-status.yml
+++ b/.github/workflows/doc-status.yml
@@ -6,10 +6,9 @@ jobs:
   see-page-build-payload:
     runs-on: ubuntu-latest
     steps:
-      - name: Set up Python 3.7
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
-          python-version: 3.7
+          python-version: 3.9
           architecture: x64
       - name: check status
         run: |
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index dbb65ac..e778325 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -9,13 +9,12 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.7]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
         with:
           persist-credentials: false
-      - name: Set up Python
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
diff --git a/.github/workflows/pr-labels.yml b/.github/workflows/pr-labels.yml
index a951045..39e039b 100644
--- a/.github/workflows/pr-labels.yml
+++ b/.github/workflows/pr-labels.yml
@@ -13,7 +13,7 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v2
       - name: Set up python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
       - name: Install requests
         run: pip install requests
       - name: Process commit and find merger responsible for labeling
diff --git a/.github/workflows/release.yml b/.github/workflows/publish.yml
similarity index 62%
rename from .github/workflows/release.yml
rename to .github/workflows/publish.yml
index d13253b..c6d9c07 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/publish.yml
@@ -1,18 +1,18 @@
-name: release
+name: publish
 
 on:
   release:
     types: [published]
 
 jobs:
-  pypi-publish:
+  pypi:
+    if: "!github.event.release.prerelease"
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.7
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
-          python-version: 3.7
+          python-version: 3.9
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2
@@ -23,30 +23,25 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install setuptools wheel twine --upgrade
-      - name: Get release tag
-        id: release_tag
-        run: |
-          echo ::set-output name=VERSION::${GITHUB_REF/refs\/tags\//}
       - name: Build and publish
         env:
           TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
           TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-          VERSION: ${{ steps.release_tag.outputs.VERSION }}
         run: |
-          BUILD_VERSION=${VERSION:1} python setup.py sdist bdist_wheel
+          echo "BUILD_VERSION=${GITHUB_REF#refs/*/}" | cut -c 2- >> $GITHUB_ENV
+          python setup.py sdist bdist_wheel
           twine check dist/*
           twine upload dist/*
 
   pypi-check:
     if: "!github.event.release.prerelease"
     runs-on: ubuntu-latest
-    needs: pypi-publish
+    needs: pypi
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.7
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
-          python-version: 3.7
+          python-version: 3.9
           architecture: x64
       - name: Install package
         run: |
@@ -54,46 +49,44 @@ jobs:
           pip install torchscan
           python -c "import torchscan; print(torchscan.__version__)"
 
-  conda-publish:
+  conda:
+    if: "!github.event.release.prerelease"
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - name: Miniconda setup
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
-          python-version: 3.7
-          auto-activate-base: true
+          python-version: 3.9
       - name: Install dependencies
-        run: |
-          conda install -y conda-build conda-verify anaconda-client
-      - name: Get release tag
-        id: release_tag
-        run: |
-          echo ::set-output name=VERSION::${GITHUB_REF/refs\/tags\//}
+        shell: bash -el {0}
+        run: conda install -y conda-build conda-verify anaconda-client
       - name: Build and publish
+        shell: bash -el {0}
         env:
           ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_TOKEN }}
-          VERSION: ${{ steps.release_tag.outputs.VERSION }}
         run: |
-          BUILD_VERSION=${VERSION:1} python setup.py sdist
+          echo "BUILD_VERSION=${GITHUB_REF#refs/*/}" | cut -c 2- >> $GITHUB_ENV
+          python setup.py sdist
           mkdir conda-dist
-          conda-build .conda/ -c pytorch --output-folder conda-dist
+          conda build .conda/ -c pytorch --output-folder conda-dist
           ls -l conda-dist/noarch/*tar.bz2
           anaconda upload conda-dist/noarch/*tar.bz2
 
   conda-check:
     if: "!github.event.release.prerelease"
     runs-on: ubuntu-latest
-    needs: conda-publish
+    needs: conda
     steps:
       - name: Miniconda setup
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
-          python-version: 3.7
+          python-version: 3.9
           auto-activate-base: true
       - name: Install package
+        shell: bash -el {0}
         run: |
           conda install -c frgfm torchscan
           python -c "import torchscan; print(torchscan.__version__)"
diff --git a/.github/workflows/pull_requests.yml b/.github/workflows/pull_requests.yml
index de64b63..2d9ac11 100644
--- a/.github/workflows/pull_requests.yml
+++ b/.github/workflows/pull_requests.yml
@@ -10,9 +10,9 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: 3.9
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2
diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml
index d16779b..57dade8 100644
--- a/.github/workflows/style.yml
+++ b/.github/workflows/style.yml
@@ -7,55 +7,33 @@ on:
     branches: main
 
 jobs:
-  flake8:
+  ruff:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.8]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
-      - name: Run flake8
+      - name: Run ruff
         run: |
-          pip install flake8
-          flake8 --version
-          flake8
-
-  isort:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ubuntu-latest]
-        python: [3.8]
-    steps:
-      - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v1
-        with:
-          python-version: ${{ matrix.python }}
-          architecture: x64
-      - name: Run isort
-        run: |
-          pip install isort
-          isort --version
-          isort .
-          if [ -n "$(git status --porcelain --untracked-files=no)" ]; then exit 1; else echo "All clear"; fi
+          pip install ruff==0.1.14
+          ruff --version
+          ruff check --diff .
 
   mypy:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.8]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
@@ -63,70 +41,50 @@ jobs:
         uses: actions/cache@v2
         with:
           path: ~/.cache/pip
-          key: ${{ runner.os }}-python-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}-mypy
+          key: ${{ runner.os }}-python-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e . --upgrade
-          pip install mypy
+          pip install -e ".[quality]" --upgrade
       - name: Run mypy
         run: |
           mypy --version
           mypy
 
-  pydocstyle:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ubuntu-latest]
-        python: [3.8]
-    steps:
-      - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python }}
-          architecture: x64
-      - name: Run pydocstyle
-        run: |
-          pip install "pydocstyle[toml]"
-          pydocstyle --version
-          pydocstyle
-
-  black:
+  ruff-format:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.8]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v2
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
-      - name: Run black
+      - name: Run ruff
         run: |
-          pip install black
-          black --version
-          black --check --diff .
+          pip install ruff==0.1.14
+          ruff --version
+          ruff format --check --diff .
 
-  bandit:
+  precommit-hooks:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.8]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python
-        uses: actions/setup-python@v2
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
-      - name: Run bandit
+      - name: Run pre-commit hooks
         run: |
-          pip install bandit[toml]
-          bandit --version
-          bandit -r . -c pyproject.toml
+          pip install pre-commit
+          git checkout -b temp
+          pre-commit install
+          pre-commit --version
+          pre-commit run --all-files
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 22e4828..a627fb4 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -12,13 +12,12 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: [3.7]
+        python: [3.9]
     steps:
       - uses: actions/checkout@v2
         with:
           persist-credentials: false
-      - name: Set up Python
-        uses: actions/setup-python@v1
+      - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python }}
           architecture: x64
@@ -32,12 +31,10 @@ jobs:
           python -m pip install --upgrade pip
           pip install -e ".[test]" --upgrade
       - name: Run unittests
-        run: |
-          coverage run -m pytest tests/
-          coverage xml
+        run: pytest --cov=torchscan --cov-report xml tests/
       - uses: actions/upload-artifact@v2
         with:
-          name: coverage-main
+          name: coverage-reports
           path: ./coverage.xml
 
   codecov-upload:
@@ -47,9 +44,11 @@ jobs:
       - uses: actions/checkout@v2
       - uses: actions/download-artifact@v2
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v1
+        uses: codecov/codecov-action@v3
         with:
+          token: ${{ secrets.CODECOV_TOKEN }}
           flags: unittests
+          directory: ./coverage-reports
           fail_ci_if_error: true
 
   headers:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b615adc..d94bf8a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,20 +1,30 @@
+default_language_version:
+    python: python3.9
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.5.0
     hooks:
-      - id: check-yaml
+      - id: check-added-large-files
+      - id: check-ast
+      - id: check-case-conflict
+      - id: check-json
+      - id: check-merge-conflict
+      - id: check-symlinks
       - id: check-toml
+      - id: check-xml
+      - id: check-yaml
+        exclude: .conda
+      - id: debug-statements
+        language_version: python3
       - id: end-of-file-fixer
+      - id: no-commit-to-branch
+        args: ['--branch', 'main']
+      - id: requirements-txt-fixer
       - id: trailing-whitespace
-  - repo: https://github.com/psf/black
-    rev: 22.3.0
-    hooks:
-    - id: black
-  - repo: https://github.com/pycqa/isort
-    rev: 5.10.1
-    hooks:
-    - id: isort
-  - repo: https://github.com/PyCQA/autoflake
-    rev: v1.7.7
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: 'v0.1.14'
     hooks:
-    - id: autoflake
+      - id: ruff
+        args:
+          - --fix
+      - id: ruff-format
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 4530014..0f7ad57 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -80,47 +80,22 @@ make test
 
 #### Code quality
 
-To run all quality checks together
+The CI will also run some sanity checks (header format, dependency consistency, etc.), which you can run as follows:
 
 ```shell
 make quality
 ```
 
-##### Lint verification
+This will read `pyproject.toml` and run:
+- lint checking, formatting ([ruff](https://docs.astral.sh/ruff/))
+- type annotation checking ([mypy](https://github.com/python/mypy))
 
-To ensure that your incoming PR complies with the lint settings, you need to install [flake8](https://flake8.pycqa.org/en/latest/), [black](https://black.readthedocs.io/en/stable/) and run the following command from the repository's root folder:
+You can apply automatic fix to most of those by running:
 
 ```shell
-flake8 ./
-black --check .
+make style
 ```
 
-##### Import order
-
-In order to ensure there is a common import order convention, run [isort](https://github.com/PyCQA/isort) as follows:
-
-```shell
-isort .
-```
-This will reorder the imports of your local files.
-
-##### Annotation typing
-
-Additionally, to catch type-related issues and have a cleaner codebase, annotation typing are expected. After installing [mypy](https://github.com/python/mypy), you can run the verifications as follows:
-
-```shell
-mypy torchscan/
-```
-
-##### Docstring style
-
-Finally, documentation being important, [pydocstyle](https://github.com/PyCQA/pydocstyle) will be checking the docstrings:
-
-```shell
-pydocstyle torchscan/
-```
-
-
 ### Submit your modifications
 
 Push your last modifications to your remote branch
diff --git a/Makefile b/Makefile
index e4ab5f2..e9885b3 100644
--- a/Makefile
+++ b/Makefile
@@ -1,22 +1,17 @@
 # this target runs checks on all files
 quality:
-	isort . -c
-	flake8
+	ruff format --check .
+	ruff check .
 	mypy
-	pydocstyle
-	black --check .
-	bandit -r . -c pyproject.toml
-	autoflake -r .
 
 # this target runs checks on all files and potentially modifies some of them
 style:
-	isort .
-	black .
-	autoflake --in-place -r .
+	ruff format .
+	ruff --fix .
 
 # Run tests for the library
 test:
-	coverage run -m pytest tests/
+	pytest --cov=torchscan tests/
 
 # Build documentation for current version
 single-docs:
diff --git a/README.md b/README.md
index 6096e22..29fb114 100644
--- a/README.md
+++ b/README.md
@@ -1,31 +1,41 @@
 <p align="center">
-  <img src="https://github.com/frgfm/torch-scan/releases/download/v0.1.1/logo_text.png" width="40%">
+  <img src="https://github.com/frgfm/torch-scan/releases/download/v0.1.1/logo_text.png" width="30%">
 </p>
 
 <p align="center">
   <a href="https://github.com/frgfm/torch-scan/actions/workflows/builds.yml">
     <img alt="CI Status" src="https://img.shields.io/github/actions/workflow/status/frgfm/torch-scan/builds.yml?branch=main&label=CI&logo=github&style=flat-square">
   </a>
-  <a href="https://frgfm.github.io/torch-scan">
-    <img src="https://img.shields.io/github/actions/workflow/status/frgfm/torch-scan/docs.yaml?branch=main&label=docs&logo=read-the-docs&style=flat-square" alt="Documentation Status">
-  </a>
-  <a href="https://codecov.io/gh/frgfm/torch-scan">
-    <img src="https://img.shields.io/codecov/c/github/frgfm/torch-scan.svg?logo=codecov&style=flat-square" alt="Test coverage percentage">
+  <a href="https://github.com/astral-sh/ruff">
+    <img src="https://img.shields.io/badge/Linter-Ruff-FCC21B?style=flat-square&logo=ruff&logoColor=white" alt="ruff">
   </a>
-  <a href="https://github.com/ambv/black">
-    <img src="https://img.shields.io/badge/code%20style-black-000000.svg?style=flat-square" alt="black">
+  <a href="https://github.com/astral-sh/ruff">
+    <img src="https://img.shields.io/badge/Formatter-Ruff-FCC21B?style=flat-square&logo=Python&logoColor=white" alt="ruff">
   </a>
   <a href="https://www.codacy.com/gh/frgfm/torch-scan/dashboard?utm_source=github.com&amp;utm_medium=referral&amp;utm_content=frgfm/torch-scan&amp;utm_campaign=Badge_Grade"><img src="https://app.codacy.com/project/badge/Grade/9dc68e8bfce34d9dbc8b44a350e9adc7"/></a>
+  <a href="https://codecov.io/gh/frgfm/torch-scan">
+    <img src="https://img.shields.io/codecov/c/github/frgfm/torch-scan.svg?logo=codecov&style=flat-square&label=Coverage" alt="Test coverage percentage">
+  </a>
 </p>
 <p align="center">
   <a href="https://pypi.org/project/torchscan/">
-    <img src="https://img.shields.io/pypi/v/torchscan.svg?logo=python&amp;logoColor=fff&amp;style=flat-square" alt="PyPi Status">
+    <img src="https://img.shields.io/pypi/v/torchscan.svg?logo=PyPI&logoColor=fff&style=flat-square&label=PyPI" alt="PyPi Version">
+  </a>
+  <a href="https://anaconda.org/frgfm/torchscan">
+    <img src="https://img.shields.io/conda/v/frgfm/torchscan.svg?logo=anaconda&label=Conda&logoColor=fff&style=flat-square" alt="Conda Version">
+  </a>
+  <img src="https://img.shields.io/pypi/pyversions/torchscan.svg?logo=Python&label=Python&logoColor=fff&style=flat-square" alt="pyversions">
+  <a href="https://github.com/frgfm/torch-scan/blob/main/LICENSE">
+    <img src="https://img.shields.io/github/license/frgfm/torch-scan.svg?label=License&logoColor=fff&style=flat-square" alt="License">
+  </a>
+</p>
+<p align="center">
+  <a href="https://frgfm.github.io/torch-scan">
+    <img src="https://img.shields.io/github/actions/workflow/status/frgfm/torch-scan/docs.yml?branch=main&label=Documentation&logo=read-the-docs&logoColor=white&style=flat-square" alt="Documentation Status">
   </a>
-  <a href="https://anaconda.org/frgfm/torchscan"> <img src="https://anaconda.org/frgfm/torchscan/badges/version.svg" /> </a>
-  <img src="https://img.shields.io/pypi/pyversions/torchscan.svg?style=flat-square" alt="pyversions">
-  <img src="https://img.shields.io/pypi/l/torchscan.svg?style=flat-square" alt="license">
 </p>
 
+
 The very useful [summary](https://www.tensorflow.org/api_docs/python/tf/keras/Model#summary) method of `tf.keras.Model` but for PyTorch, with more useful information.
 
 
@@ -104,7 +114,7 @@ which will add the layer's receptive field (relatively to the last convolutional
 
 ## Setup
 
-Python 3.6 (or higher) and [pip](https://pip.pypa.io/en/stable/)/[conda](https://docs.conda.io/en/latest/miniconda.html) are required to install Torchscan.
+Python 3.8 (or newer) and [pip](https://pip.pypa.io/en/stable/)/[conda](https://docs.conda.io/en/latest/miniconda.html) are required to install Torchscan.
 
 ### Stable release
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 819b70b..ba0eff5 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -15,12 +15,11 @@
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
-import os
 import sys
-
-sys.path.insert(0, os.path.abspath("../.."))
 from datetime import datetime
+from pathlib import Path
 
+sys.path.insert(0, Path().cwd().parent.parent)
 import torchscan
 
 # -- Project information -----------------------------------------------------
@@ -106,10 +105,10 @@
 # so a file named "default.css" will overwrite the builtin "default.css".
 html_static_path = ["_static"]
 
+
 # Add googleanalytics id
 # ref: https://github.com/orenhecht/googleanalytics/blob/master/sphinxcontrib/googleanalytics.py
 def add_ga_javascript(app, pagename, templatename, context, doctree):
-
     metatags = context.get("metatags", "")
     metatags += """
     <!-- Global site tag (gtag.js) - Google Analytics -->
@@ -120,9 +119,7 @@ def add_ga_javascript(app, pagename, templatename, context, doctree):
   gtag('js', new Date());
   gtag('config', '{0}');
 </script>
-    """.format(
-        app.config.googleanalytics_id
-    )
+    """.format(app.config.googleanalytics_id)
     context["metatags"] = metatags
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 3d2deba..fba5175 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ authors = [
     {name = "François-Guillaume Fernandez", email = "fg-feedback@protonmail.com"}
 ]
 readme = "README.md"
-requires-python = ">=3.6,<4"
+requires-python = ">=3.8,<4"
 license = {file = "LICENSE"}
 keywords = ["pytorch", "deep learning", "summary", "memory", "ram"]
 classifiers = [
@@ -20,33 +20,29 @@ classifiers = [
     "Natural Language :: English",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.6",
-    "Programming Language :: Python :: 3.7",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Mathematics",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dynamic = ["version"]
 dependencies = [
-    "torch>=1.5.0,<2.0.0",
+    "torch>=2.0.0,<3.0.0",
 ]
 
 [project.optional-dependencies]
 test = [
-    "pytest>=5.3.2",
-    "coverage[toml]>=4.5.4",
+    "pytest>=7.3.2",
+    "pytest-cov>=3.0.0,<5.0.0",
+    "pytest-pretty>=1.0.0,<2.0.0",
 ]
 quality = [
-    "flake8>=3.9.0",
-    "isort>=5.7.0",
-    "mypy>=0.812",
-    "pydocstyle[toml]>=6.0.0",
-    "black>=22.1,<23.0",
-    "autoflake>=1.5.0,<2.0.0",
-    "bandit[toml]>=1.7.0,<1.8.0",
-    "pre-commit>=2.17.0,<3.0.0",
+    "ruff==0.1.14",
+    "mypy==1.8.0",
+    "pre-commit>=3.0.0,<4.0.0",
 ]
 docs = [
     "sphinx>=3.0.0,!=3.5.0",
@@ -59,17 +55,13 @@ docs = [
 ]
 dev = [
     # test
-    "pytest>=5.3.2",
-    "coverage[toml]>=4.5.4",
+    "pytest>=7.3.2",
+    "pytest-cov>=3.0.0,<5.0.0",
+    "pytest-pretty>=1.0.0,<2.0.0",
     # style
-    "flake8>=3.9.0",
-    "isort>=5.7.0",
-    "mypy>=0.812",
-    "pydocstyle[toml]>=6.0.0",
-    "black>=22.1,<23.0",
-    "autoflake>=1.5.0,<2.0.0",
-    "bandit[toml]>=1.7.0,<1.8.0",
-    "pre-commit>=2.17.0,<3.0.0",
+    "ruff==0.1.14",
+    "mypy==1.8.0",
+    "pre-commit>=3.0.0,<4.0.0",
     # docs
     "sphinx>=3.0.0,!=3.5.0",
     "furo>=2022.3.4",
@@ -90,8 +82,91 @@ zip-safe = true
 [tool.setuptools.packages.find]
 exclude = ["docs*", "scripts*", "tests*"]
 
+[tool.pytest.ini_options]
+testpaths = ["torchscan/"]
+
+[tool.coverage.run]
+source = ["torchscan/"]
+
+[tool.ruff]
+select = [
+    "E",  # pycodestyle errors
+    "W",  # pycodestyle warnings
+    "D101", "D103",  # pydocstyle missing docstring in public function/class
+    "D201","D202","D207","D208","D214","D215","D300","D301","D417", "D419",  # pydocstyle
+    "F",  # pyflakes
+    "I",  # isort
+    "C4",  # flake8-comprehensions
+    "B",  # flake8-bugbear
+    "CPY001",  # flake8-copyright
+    "ISC",  # flake8-implicit-str-concat
+    "PYI",  # flake8-pyi
+    "NPY",  # numpy
+    "PERF",  # perflint
+    "RUF",  # ruff specific
+    "PTH",  # flake8-use-pathlib
+    "S",  # flake8-bandit
+    "N",  # pep8-naming
+    "T10",  # flake8-debugger
+    "T20",  # flake8-print
+    "PT",  # flake8-pytest-style
+    "LOG",  # flake8-logging
+    "SIM",  # flake8-simplify
+    "YTT",  # flake8-2020
+    "ANN",  # flake8-annotations
+    "ASYNC",  # flake8-async
+    "BLE",  # flake8-blind-except
+    "A",  # flake8-builtins
+    "ICN",  # flake8-import-conventions
+    "PIE",  # flake8-pie
+    "ARG",  # flake8-unused-arguments
+    "FURB",  # refurb
+]
+ignore = [
+    "E501",  # line too long, handled by black
+    "B008",  # do not perform function calls in argument defaults
+    "B904",  # raise from
+    "C901",  # too complex
+    "F403",  # star imports
+    "E731",  # lambda assignment
+    "C416",  # list comprehension to list()
+    "ANN101",  # missing type annotations on self
+    "ANN102",  # missing type annotations on cls
+    "ANN002",  # missing type annotations on *args
+    "ANN003",  # missing type annotations on **kwargs
+    "COM812",  # trailing comma missing
+    "N812",  # lowercase imported as non-lowercase
+    "ISC001",  # implicit string concatenation (handled by format)
+    "ANN401",  # Dynamically typed expressions (typing.Any) are disallowed
+]
+exclude = [".git"]
+line-length = 120
+target-version = "py39"
+preview = true
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+
+[tool.ruff.per-file-ignores]
+"**/__init__.py" = ["I001", "F401", "CPY001"]
+"scripts/**.py" = ["D", "T201", "N812", "S101", "ANN"]
+".github/**.py" = ["D", "T201", "S602", "S101", "ANN"]
+"docs/**.py" = ["E402", "D103", "ANN", "A001", "ARG001"]
+"tests/**.py" = ["D101", "D103", "CPY001", "S101", "PT011", "ANN"]
+"demo/**.py" = ["D103", "ANN"]
+"setup.py" = ["T201"]
+"torchscan/process/memory.py" = ["S60"]
+
+[tool.ruff.flake8-quotes]
+docstring-quotes = "double"
+
+[tool.ruff.isort]
+known-first-party = ["torchscan", "app"]
+known-third-party = ["torch", "torchvision"]
 
 [tool.mypy]
+python_version = "3.9"
 files = "torchscan/"
 show_error_codes = true
 pretty = true
@@ -101,29 +176,5 @@ no_implicit_optional = true
 disallow_untyped_calls = true
 check_untyped_defs = true
 implicit_reexport = false
-
-[tool.isort]
-line_length = 120
-src_paths = ["torchscan", "tests", "scripts", "docs", ".github"]
-skip_glob = "**/__init__.py"
-known_third_party = ["torch", "torchvision"]
-
-[tool.pydocstyle]
-select = "D300,D301,D417"
-match = ".*\\.py"
-
-[tool.coverage.run]
-source = ["torchscan"]
-
-[tool.black]
-line-length = 120
-target-version = ['py38']
-
-[tool.autoflake]
-remove-unused-variables = true
-remove-all-unused-imports = true
-ignore-init-module-imports = true
-
-[tool.bandit]
-exclude_dirs = [".github/collect_env.py", "torchscan/process/memory.py"]
-skips = ["B101"]
+disallow_untyped_defs = true
+explicit_package_bases = true
diff --git a/scripts/benchmark.py b/scripts/benchmark.py
index 7ad011a..570f129 100644
--- a/scripts/benchmark.py
+++ b/scripts/benchmark.py
@@ -52,7 +52,6 @@
 
 
 def main():
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     margin = 4
diff --git a/setup.py b/setup.py
index 7dd2bdb..a1db9a1 100644
--- a/setup.py
+++ b/setup.py
@@ -14,12 +14,11 @@
 
 
 if __name__ == "__main__":
-
     print(f"Building wheel {PKG_NAME}-{VERSION}")
 
     # Dynamically set the __version__ attribute
     cwd = Path(__file__).parent.absolute()
-    with open(cwd.joinpath("torchscan", "version.py"), "w", encoding="utf-8") as f:
+    with cwd.joinpath("torchscan", "version.py").open("w", encoding="utf-8") as f:
         f.write(f"__version__ = '{VERSION}'\n")
 
     setup(name=PKG_NAME, version=VERSION)
diff --git a/tests/test_crawler.py b/tests/test_crawler.py
index 0a73247..5f160d4 100644
--- a/tests/test_crawler.py
+++ b/tests/test_crawler.py
@@ -24,7 +24,6 @@ def tag_name(mod, name):
 
 
 def test_crawl_module():
-
     mod = nn.Conv2d(3, 8, 3)
 
     res = crawler.crawl_module(mod, (3, 32, 32))
@@ -34,7 +33,6 @@ def test_crawl_module():
 
 
 def test_summary():
-
     mod = nn.Conv2d(3, 8, 3)
 
     # Redirect stdout with StringIO object
@@ -67,13 +65,11 @@ def test_summary():
         crawler.summary(mod, (3, 32, 32), max_depth=1)
 
     mod = nn.Sequential(
-        OrderedDict(
-            [
-                ("features", nn.Sequential(nn.Conv2d(3, 8, 3), nn.ReLU(inplace=True))),
-                ("pool", nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(1))),
-                ("classifier", nn.Linear(8, 1)),
-            ]
-        )
+        OrderedDict([
+            ("features", nn.Sequential(nn.Conv2d(3, 8, 3), nn.ReLU(inplace=True))),
+            ("pool", nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(1))),
+            ("classifier", nn.Linear(8, 1)),
+        ])
     )
 
     captured_output = io.StringIO()
diff --git a/tests/test_modules.py b/tests/test_modules.py
index 40fccb5..349880d 100644
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -16,38 +16,38 @@ def test_module_flops_warning():
 
 
 @pytest.mark.parametrize(
-    "mod, input_shape, output_shape, expected_val",
+    ("mod", "input_shape", "output_shape", "expected_val"),
     [
         # Check for unknown module that it returns 0 and throws a warning
-        [MyModule(), (1,), (1,), 0],
+        (MyModule(), (1,), (1,), 0),
         # Fully-connected
-        [nn.Linear(8, 4), (1, 8), (1, 4), 4 * (2 * 8 - 1) + 4],
-        [nn.Linear(8, 4, bias=False), (1, 8), (1, 4), 4 * (2 * 8 - 1)],
-        [nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 2 * (4 * (2 * 8 - 1) + 4)],
+        (nn.Linear(8, 4), (1, 8), (1, 4), 4 * (2 * 8 - 1) + 4),
+        (nn.Linear(8, 4, bias=False), (1, 8), (1, 4), 4 * (2 * 8 - 1)),
+        (nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 2 * (4 * (2 * 8 - 1) + 4)),
         # Activations
-        [nn.Identity(), (1, 8), (1, 8), 0],
-        [nn.Flatten(), (1, 8), (1, 8), 0],
-        [nn.ReLU(), (1, 8), (1, 8), 8],
-        [nn.ELU(), (1, 8), (1, 8), 48],
-        [nn.LeakyReLU(), (1, 8), (1, 8), 32],
-        [nn.ReLU6(), (1, 8), (1, 8), 16],
-        [nn.Tanh(), (1, 8), (1, 8), 48],
-        [nn.Sigmoid(), (1, 8), (1, 8), 32],
+        (nn.Identity(), (1, 8), (1, 8), 0),
+        (nn.Flatten(), (1, 8), (1, 8), 0),
+        (nn.ReLU(), (1, 8), (1, 8), 8),
+        (nn.ELU(), (1, 8), (1, 8), 48),
+        (nn.LeakyReLU(), (1, 8), (1, 8), 32),
+        (nn.ReLU6(), (1, 8), (1, 8), 16),
+        (nn.Tanh(), (1, 8), (1, 8), 48),
+        (nn.Sigmoid(), (1, 8), (1, 8), 32),
         # BN
-        [nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 144 + 32 + 32 * 3 + 48],
+        (nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 144 + 32 + 32 * 3 + 48),
         # Pooling
-        [nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
-        [nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AdaptiveAvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
-        [nn.AdaptiveAvgPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
+        (nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
+        (nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AdaptiveAvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
+        (nn.AdaptiveAvgPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
         # Dropout
-        [nn.Dropout(), (1, 8), (1, 8), 8],
-        [nn.Dropout(p=0), (1, 8), (1, 8), 0],
+        (nn.Dropout(), (1, 8), (1, 8), 8),
+        (nn.Dropout(p=0), (1, 8), (1, 8), 0),
         # Conv
-        [nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 388800],
-        [nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 499408],
+        (nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 388800),
+        (nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 499408),
     ],
 )
 def test_module_flops(mod, input_shape, output_shape, expected_val):
@@ -67,33 +67,32 @@ def test_module_macs_warning():
 
 
 @pytest.mark.parametrize(
-    "mod, input_shape, output_shape, expected_val",
+    ("mod", "input_shape", "output_shape", "expected_val"),
     [
         # Check for unknown module that it returns 0 and throws a warning
-        [MyModule(), (1,), (1,), 0],
+        (MyModule(), (1,), (1,), 0),
         # Fully-connected
-        [nn.Linear(8, 4), (1, 8), (1, 4), 8 * 4],
-        [nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 8 * 4 * 2],
+        (nn.Linear(8, 4), (1, 8), (1, 4), 8 * 4),
+        (nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 8 * 4 * 2),
         # Activations
-        [nn.ReLU(), (1, 8), (1, 8), 0],
+        (nn.ReLU(), (1, 8), (1, 8), 0),
         # BN
-        [nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 64 + 24 + 56 + 32],
+        (nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 64 + 24 + 56 + 32),
         # Pooling
-        [nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
-        [nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32],
-        [nn.AdaptiveAvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
-        [nn.AdaptiveAvgPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32],
+        (nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
+        (nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 3 * 32),
+        (nn.AdaptiveAvgPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
+        (nn.AdaptiveAvgPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 5 * 32),
         # Dropout
-        [nn.Dropout(), (1, 8), (1, 8), 0],
+        (nn.Dropout(), (1, 8), (1, 8), 0),
         # Conv
-        [nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 194400],
-        [nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 249704],
+        (nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 194400),
+        (nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 249704),
     ],
 )
 def test_module_macs(mod, input_shape, output_shape, expected_val):
-
     assert modules.module_macs(mod, torch.zeros(input_shape), torch.zeros(output_shape)) == expected_val
 
 
@@ -103,37 +102,36 @@ def test_module_dmas_warning():
 
 
 @pytest.mark.parametrize(
-    "mod, input_shape, output_shape, expected_val",
+    ("mod", "input_shape", "output_shape", "expected_val"),
     [
         # Check for unknown module that it returns 0 and throws a warning
-        [MyModule(), (1,), (1,), 0],
+        (MyModule(), (1,), (1,), 0),
         # Fully-connected
-        [nn.Linear(8, 4), (1, 8), (1, 4), 4 * (8 + 1) + 8 + 4],
-        [nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 4 * (8 + 1) + 2 * (8 + 4)],
+        (nn.Linear(8, 4), (1, 8), (1, 4), 4 * (8 + 1) + 8 + 4),
+        (nn.Linear(8, 4), (1, 2, 8), (1, 2, 4), 4 * (8 + 1) + 2 * (8 + 4)),
         # Activations
-        [nn.Identity(), (1, 8), (1, 8), 8],
-        [nn.Flatten(), (1, 8), (1, 8), 16],
-        [nn.ReLU(), (1, 8), (1, 8), 8 * 2],
-        [nn.ReLU(inplace=True), (1, 8), (1, 8), 8],
-        [nn.ELU(), (1, 8), (1, 8), 17],
-        [nn.Tanh(), (1, 8), (1, 8), 24],
-        [nn.Sigmoid(), (1, 8), (1, 8), 16],
+        (nn.Identity(), (1, 8), (1, 8), 8),
+        (nn.Flatten(), (1, 8), (1, 8), 16),
+        (nn.ReLU(), (1, 8), (1, 8), 8 * 2),
+        (nn.ReLU(inplace=True), (1, 8), (1, 8), 8),
+        (nn.ELU(), (1, 8), (1, 8), 17),
+        (nn.Tanh(), (1, 8), (1, 8), 24),
+        (nn.Sigmoid(), (1, 8), (1, 8), 16),
         # BN
-        [nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 32 + 17 + 16 + 1 + 17 + 32],
+        (nn.BatchNorm1d(8), (1, 8, 4), (1, 8, 4), 32 + 17 + 16 + 1 + 17 + 32),
         # Pooling
-        [nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32],
-        [nn.MaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32],
-        [nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32],
-        [nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32],
+        (nn.MaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32),
+        (nn.MaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32),
+        (nn.AdaptiveMaxPool2d((2, 2)), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32),
+        (nn.AdaptiveMaxPool2d(2), (1, 8, 4, 4), (1, 8, 2, 2), 4 * 32 + 32),
         # Dropout
-        [nn.Dropout(), (1, 8), (1, 8), 17],
+        (nn.Dropout(), (1, 8), (1, 8), 17),
         # Conv
-        [nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 201824],
-        [nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 259178],
+        (nn.Conv2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 30, 30), 201824),
+        (nn.ConvTranspose2d(3, 8, 3), (1, 3, 32, 32), (1, 8, 34, 34), 259178),
     ],
 )
 def test_module_dmas(mod, input_shape, output_shape, expected_val):
-
     assert modules.module_dmas(mod, torch.zeros(input_shape), torch.zeros(output_shape)) == expected_val
 
 
@@ -169,7 +167,7 @@ def test_module_dmas(mod, input_shape, output_shape, expected_val):
 #     self.assertEqual(modules.module_rf(nn.BatchNorm1d(8), torch.zeros((1, 8, 4)), torch.zeros((1, 8, 4))),
 #                      (1, 1, 0))
 
-#     # Pooling
+#     # Pooling
 #     self.assertEqual(modules.module_rf(nn.MaxPool2d((2, 2)),
 #                                        torch.zeros((1, 8, 4, 4)), torch.zeros((1, 8, 2, 2))),
 #                      (2, 2, 0))
diff --git a/tests/test_process.py b/tests/test_process.py
index b40371a..319cca1 100644
--- a/tests/test_process.py
+++ b/tests/test_process.py
@@ -6,7 +6,6 @@
 
 
 def test_get_process_gpu_ram():
-
     if torch.cuda.is_initialized:
         assert process.get_process_gpu_ram(os.getpid()) >= 0
     else:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 03d263a..3edafdc 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -11,7 +11,6 @@ def test_format_name():
 
 
 def test_wrap_string():
-
     example = ".".join(["a" for _ in range(10)])
     max_len = 10
     wrap = "[...]"
@@ -24,13 +23,13 @@ def test_wrap_string():
 
 
 @pytest.mark.parametrize(
-    "input_val, num_val, unit",
+    ("input_val", "num_val", "unit"),
     [
-        [3e14, 300, "T"],
-        [3e10, 30, "G"],
-        [3e7, 30, "M"],
-        [15e3, 15, "k"],
-        [500, 500, ""],
+        (3e14, 300, "T"),
+        (3e10, 30, "G"),
+        (3e7, 30, "M"),
+        (15e3, 15, "k"),
+        (500, 500, ""),
     ],
 )
 def test_unit_scale(input_val, num_val, unit):
diff --git a/torchscan/__init__.py b/torchscan/__init__.py
index 8aac9ef..f5eb582 100644
--- a/torchscan/__init__.py
+++ b/torchscan/__init__.py
@@ -1,7 +1,6 @@
+from contextlib import suppress
 from torchscan import modules, process, utils
 from torchscan.crawler import *
 
-try:
-    from .version import __version__  # noqa: F401
-except ImportError:
-    pass
+with suppress(ImportError):
+    from .version import __version__
diff --git a/torchscan/crawler.py b/torchscan/crawler.py
index 2dc52ea..c98dc38 100644
--- a/torchscan/crawler.py
+++ b/torchscan/crawler.py
@@ -24,7 +24,6 @@ def apply(module: Module, fn: Callable[[Module, str], None], name: Optional[str]
         fn: function to apply to each module
         name: name of the current module
     """
-
     if name is None:
         name = module.__class__.__name__.lower()
     fn(module, name)
@@ -51,7 +50,6 @@ def crawl_module(
     Returns:
         layer and overhead information
     """
-
     # Get device and data types from model
     p = next(module.parameters())
     device = p.device
@@ -118,28 +116,26 @@ def _pre_hook(module: Module, inp: torch.Tensor) -> None:
                 else:
                     call_idxs[id(module)].append(len(info))
 
-                info.append(
-                    dict(
-                        name=name.rpartition(".")[-1],
-                        depth=len(name.split(".")) - 1,
-                        type=module.__class__.__name__,
-                        input_shape=(-1, *inp[0][0].shape[1:]),
-                        output_shape=None,
-                        grad_params=grad_params,
-                        nograd_params=nograd_params,
-                        param_size=param_size,
-                        num_buffers=num_buffers,
-                        buffer_size=buffer_size,
-                        flops=0,
-                        macs=0,
-                        dmas=0,
-                        rf=1,
-                        s=1,
-                        p=0,
-                        is_shared=is_shared,
-                        is_leaf=not any(module.children()),
-                    )
-                )
+                info.append({
+                    "name": name.rpartition(".")[-1],
+                    "depth": len(name.split(".")) - 1,
+                    "type": module.__class__.__name__,
+                    "input_shape": (-1, *inp[0][0].shape[1:]),
+                    "output_shape": None,
+                    "grad_params": grad_params,
+                    "nograd_params": nograd_params,
+                    "param_size": param_size,
+                    "num_buffers": num_buffers,
+                    "buffer_size": buffer_size,
+                    "flops": 0,
+                    "macs": 0,
+                    "dmas": 0,
+                    "rf": 1,
+                    "s": 1,
+                    "p": 0,
+                    "is_shared": is_shared,
+                    "is_leaf": not any(module.children()),
+                })
                 # Mark the next hook for execution
                 pre_hook_tracker[id(module)]["target"] += 1
                 # Current pass already used one of the hooks
@@ -152,7 +148,6 @@ def _pre_hook(module: Module, inp: torch.Tensor) -> None:
 
         def _fwd_hook(module: Module, inputs: Tuple[torch.Tensor, ...], out: torch.Tensor) -> None:
             """Post-forward hook"""
-
             # Check that another hook has not been triggered at this forward stage
             if not post_hook_tracker[id(module)]["is_used"] and (
                 post_hook_tracker[id(module)]["target"] == post_hook_tracker[id(module)]["current"]
@@ -199,11 +194,11 @@ def _fwd_hook(module: Module, inputs: Tuple[torch.Tensor, ...], out: torch.Tenso
                 post_hook_tracker[id(module)]["current"] = 0
                 post_hook_tracker[id(module)]["is_used"] = False
 
-        pre_fw_handles.append(module.register_forward_pre_hook(_pre_hook))
+        pre_fw_handles.append(module.register_forward_pre_hook(_pre_hook))  # type: ignore[arg-type]
         post_fw_handles.append(module.register_forward_hook(_fwd_hook))
         # Handle modules that are used multiple times (with several hooks)
-        pre_hook_tracker[id(module)] = dict(current=0, target=0, is_used=False)
-        post_hook_tracker[id(module)] = dict(current=0, target=0, is_used=False)
+        pre_hook_tracker[id(module)] = {"current": 0, "target": 0, "is_used": False}
+        post_hook_tracker[id(module)] = {"current": 0, "target": 0, "is_used": False}
 
     # Hook model
     info: List[Dict[str, Any]] = []
@@ -250,23 +245,23 @@ def _fwd_hook(module: Module, inputs: Tuple[torch.Tensor, ...], out: torch.Tenso
         info[fw_idx]["s"] = _s
         info[fw_idx]["p"] = _p
 
-    return dict(
-        overheads=dict(
-            cuda=dict(
-                pre=cuda_overhead,
-                fwd=get_process_gpu_ram(os.getpid()) - reserved_ram,
-            ),
-            framework=dict(pre=framework_overhead, fwd=diff_ram),
-        ),
-        layers=info,
-        overall=dict(
-            grad_params=grad_params,
-            nograd_params=nograd_params,
-            param_size=param_size,
-            num_buffers=num_buffers,
-            buffer_size=buffer_size,
-        ),
-    )
+    return {
+        "overheads": {
+            "cuda": {
+                "pre": cuda_overhead,
+                "fwd": get_process_gpu_ram(os.getpid()) - reserved_ram,
+            },
+            "framework": {"pre": framework_overhead, "fwd": diff_ram},
+        },
+        "layers": info,
+        "overall": {
+            "grad_params": grad_params,
+            "nograd_params": nograd_params,
+            "param_size": param_size,
+            "num_buffers": num_buffers,
+            "buffer_size": buffer_size,
+        },
+    }
 
 
 def summary(
@@ -292,11 +287,10 @@ def summary(
         receptive_field: whether receptive field estimation should be performed
         effective_rf_stats: if `receptive_field` is True, displays effective stride and padding
     """
-
     # Get the summary dict
     module_info = crawl_module(module, input_shape)
     # Aggregate until max_depth
     if isinstance(max_depth, int):
         module_info = aggregate_info(module_info, max_depth)
     # Format it and print it
-    print(format_info(module_info, wrap_mode, receptive_field, effective_rf_stats))
+    print(format_info(module_info, wrap_mode, receptive_field, effective_rf_stats))  # noqa T201
diff --git a/torchscan/modules/__init__.py b/torchscan/modules/__init__.py
index 582290a..7676278 100644
--- a/torchscan/modules/__init__.py
+++ b/torchscan/modules/__init__.py
@@ -1,4 +1,4 @@
 from .flops import *
-from .macs import *  # type: ignore[assignment]
+from .macs import *
 from .memory import *
 from .receptive import *
diff --git a/torchscan/modules/flops.py b/torchscan/modules/flops.py
index 9550b8c..296ccf8 100644
--- a/torchscan/modules/flops.py
+++ b/torchscan/modules/flops.py
@@ -28,7 +28,6 @@ def module_flops(module: Module, inputs: Tuple[Tensor, ...], out: Tensor) -> int
     Returns:
         number of FLOPs
     """
-
     if isinstance(module, (nn.Identity, nn.Flatten)):
         return 0
     elif isinstance(module, nn.Linear):
@@ -64,13 +63,12 @@ def module_flops(module: Module, inputs: Tuple[Tensor, ...], out: Tensor) -> int
     elif isinstance(module, nn.Transformer):
         return flops_transformer(module, inputs)
     else:
-        warnings.warn(f"Module type not supported: {module.__class__.__name__}")
+        warnings.warn(f"Module type not supported: {module.__class__.__name__}", stacklevel=1)
         return 0
 
 
 def flops_linear(module: nn.Linear, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.Linear`"""
-
     # batch size * out_chan * in_chan
     num_out_feats = module.out_features * reduce(mul, inputs[0].shape[:-1])
     mm_flops = num_out_feats * (2 * module.in_features - 1)
@@ -79,51 +77,44 @@ def flops_linear(module: nn.Linear, inputs: Tuple[Tensor, ...]) -> int:
     return mm_flops + bias_flops
 
 
-def flops_sigmoid(module: nn.Sigmoid, inputs: Tuple[Tensor, ...]) -> int:
+def flops_sigmoid(_: nn.Sigmoid, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.Sigmoid`"""
-
     # For each element, mul by -1, exp it, add 1, div
     return inputs[0].numel() * 4
 
 
-def flops_relu(module: nn.ReLU, inputs: Tuple[Tensor, ...]) -> int:
+def flops_relu(_: nn.ReLU, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.ReLU`"""
-
     # Each element is compared to 0
     return inputs[0].numel()
 
 
-def flops_elu(module: nn.ELU, inputs: Tuple[Tensor, ...]) -> int:
+def flops_elu(_: nn.ELU, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.ELU`"""
-
     # For each element, compare it to 0, exp it, sub 1, mul by alpha, compare it to 0 and sum both
     return inputs[0].numel() * 6
 
 
-def flops_leakyrelu(module: nn.LeakyReLU, inputs: Tuple[Tensor, ...]) -> int:
+def flops_leakyrelu(_: nn.LeakyReLU, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.LeakyReLU`"""
-
     # For each element, compare it to 0 (max), compare it to 0 (min), mul by slope and sum both
     return inputs[0].numel() * 4
 
 
-def flops_relu6(module: nn.ReLU6, inputs: Tuple[Tensor, ...]) -> int:
+def flops_relu6(_: nn.ReLU6, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.ReLU6`"""
-
     # For each element, compare it to 0 (max), compare it to 0 (min), mul by slope and sum both
     return inputs[0].numel() * 2
 
 
-def flops_tanh(module: nn.Tanh, inputs: Tuple[Tensor, ...]) -> int:
+def flops_tanh(_: nn.Tanh, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.Tanh`"""
-
     # For each element, exp it, mul by -1 and exp it, divide the sub by the add
     return inputs[0].numel() * 6
 
 
 def flops_dropout(module: nn.Dropout, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.Dropout`"""
-
     if module.p > 0:
         # Sample a random number for each input element
         return inputs[0].numel()
@@ -133,7 +124,6 @@ def flops_dropout(module: nn.Dropout, inputs: Tuple[Tensor, ...]) -> int:
 
 def flops_convtransposend(module: _ConvTransposeNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.conv._ConvTranposeNd`"""
-
     # Padding (# cf. https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py#L496-L532)
     # Define min and max sizes
     padding_flops = len(module.kernel_size) * 8
@@ -146,7 +136,6 @@ def flops_convtransposend(module: _ConvTransposeNd, inputs: Tuple[Tensor, ...],
 
 def flops_convnd(module: _ConvNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.conv._ConvNd`"""
-
     # For each position, # mult = kernel size, # adds = kernel size - 1
     window_flops_per_chan = 2 * reduce(mul, module.kernel_size) - 1
     # Connections to input channels is controlled by the group parameter
@@ -163,7 +152,6 @@ def flops_convnd(module: _ConvNd, inputs: Tuple[Tensor, ...], out: Tensor) -> in
 
 def flops_bn(module: _BatchNorm, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.modules.batchnorm._BatchNorm`"""
-
     # for each channel, add eps and running_var, sqrt it
     norm_ops = module.num_features * 2
     # For each element, sub running_mean, div by denom
@@ -189,9 +177,8 @@ def flops_bn(module: _BatchNorm, inputs: Tuple[Tensor, ...]) -> int:
     return bn_flops + tracking_flops
 
 
-def flops_maxpool(module: _MaxPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
+def flops_maxpool(module: _MaxPoolNd, _: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.pooling._MaxPoolNd`"""
-
     k_size = reduce(mul, module.kernel_size) if isinstance(module.kernel_size, tuple) else module.kernel_size
 
     # for each spatial output element, check max element in kernel scope
@@ -200,16 +187,14 @@ def flops_maxpool(module: _MaxPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -
 
 def flops_avgpool(module: _AvgPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.pooling._AvgPoolNd`"""
-
     k_size = reduce(mul, module.kernel_size) if isinstance(module.kernel_size, tuple) else module.kernel_size
 
     # for each spatial output element, sum elements in kernel scope and div by kernel size
     return out.numel() * (k_size - 1 + inputs[0].ndim - 2)
 
 
-def flops_adaptive_maxpool(module: _AdaptiveMaxPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
+def flops_adaptive_maxpool(_: _AdaptiveMaxPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.pooling._AdaptiveMaxPoolNd`"""
-
     # Approximate kernel_size using ratio of spatial shapes between input and output
     kernel_size = tuple(
         i_size // o_size if (i_size % o_size) == 0 else i_size - o_size * (i_size // o_size) + 1
@@ -220,9 +205,8 @@ def flops_adaptive_maxpool(module: _AdaptiveMaxPoolNd, inputs: Tuple[Tensor, ...
     return out.numel() * (reduce(mul, kernel_size) - 1)
 
 
-def flops_adaptive_avgpool(module: _AdaptiveAvgPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
+def flops_adaptive_avgpool(_: _AdaptiveAvgPoolNd, inputs: Tuple[Tensor, ...], out: Tensor) -> int:
     """FLOPs estimation for `torch.nn.modules.pooling._AdaptiveAvgPoolNd`"""
-
     # Approximate kernel_size using ratio of spatial shapes between input and output
     kernel_size = tuple(
         i_size // o_size if (i_size % o_size) == 0 else i_size - o_size * (i_size // o_size) + 1
@@ -235,7 +219,6 @@ def flops_adaptive_avgpool(module: _AdaptiveAvgPoolNd, inputs: Tuple[Tensor, ...
 
 def flops_layernorm(module: nn.LayerNorm, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.modules.batchnorm._BatchNorm`"""
-
     # Compute current mean
     norm_ops = reduce(mul, module.normalized_shape) * inputs[0].shape[: -len(module.normalized_shape)].numel()
     # current var (sub the mean, square it, sum them, divide by remaining shape)
@@ -252,7 +235,6 @@ def flops_layernorm(module: nn.LayerNorm, inputs: Tuple[Tensor, ...]) -> int:
 
 def flops_mha(module: nn.MultiheadAttention, inputs: Tuple[Tensor, ...]) -> int:
     """FLOPs estimation for `torch.nn.MultiheadAttention`"""
-
     # Input projection
     q, k, _ = inputs[:3]
     batch_size = q.shape[1]
diff --git a/torchscan/modules/macs.py b/torchscan/modules/macs.py
index cd91666..8b596f8 100644
--- a/torchscan/modules/macs.py
+++ b/torchscan/modules/macs.py
@@ -47,13 +47,12 @@ def module_macs(module: Module, inp: Tensor, out: Tensor) -> int:
     elif isinstance(module, nn.Dropout):
         return 0
     else:
-        warnings.warn(f"Module type not supported: {module.__class__.__name__}")
+        warnings.warn(f"Module type not supported: {module.__class__.__name__}", stacklevel=1)
         return 0
 
 
-def macs_linear(module: nn.Linear, inp: Tensor, out: Tensor) -> int:
+def macs_linear(module: nn.Linear, _: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.Linear`"""
-
     # batch size * out_chan * macs_per_elt (bias already counted in accumulation)
     mm_mac = module.in_features * reduce(mul, out.shape)
 
@@ -62,7 +61,6 @@ def macs_linear(module: nn.Linear, inp: Tensor, out: Tensor) -> int:
 
 def macs_convtransposend(module: _ConvTransposeNd, inp: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.conv._ConvTransposeNd`"""
-
     # Padding (# cf. https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py#L496-L532)
     # Define min and max sizes, then subtract them
     padding_macs = len(module.kernel_size) * 4
@@ -75,7 +73,6 @@ def macs_convtransposend(module: _ConvTransposeNd, inp: Tensor, out: Tensor) ->
 
 def macs_convnd(module: _ConvNd, inp: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.conv._ConvNd`"""
-
     # For each position, # mult = kernel size, # adds = kernel size - 1
     window_macs_per_chan = reduce(mul, module.kernel_size)
     # Connections to input channels is controlled by the group parameter
@@ -88,9 +85,8 @@ def macs_convnd(module: _ConvNd, inp: Tensor, out: Tensor) -> int:
     return conv_mac
 
 
-def macs_bn(module: _BatchNorm, inp: Tensor, out: Tensor) -> int:
+def macs_bn(module: _BatchNorm, inp: Tensor, _: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.batchnorm._BatchNorm`"""
-
     # sub mean, div by denom
     norm_mac = 1
     # mul by gamma, add beta
@@ -116,9 +112,8 @@ def macs_bn(module: _BatchNorm, inp: Tensor, out: Tensor) -> int:
     return bn_mac + tracking_mac
 
 
-def macs_maxpool(module: _MaxPoolNd, inp: Tensor, out: Tensor) -> int:
+def macs_maxpool(module: _MaxPoolNd, _: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.pooling._MaxPoolNd`"""
-
     k_size = reduce(mul, module.kernel_size) if isinstance(module.kernel_size, tuple) else module.kernel_size
 
     # for each spatial output element, check max element in kernel scope
@@ -127,16 +122,14 @@ def macs_maxpool(module: _MaxPoolNd, inp: Tensor, out: Tensor) -> int:
 
 def macs_avgpool(module: _AvgPoolNd, inp: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.pooling._AvgPoolNd`"""
-
     k_size = reduce(mul, module.kernel_size) if isinstance(module.kernel_size, tuple) else module.kernel_size
 
     # for each spatial output element, sum elements in kernel scope and div by kernel size
     return out.numel() * (k_size - 1 + inp.ndim - 2)
 
 
-def macs_adaptive_maxpool(module: _AdaptiveMaxPoolNd, inp: Tensor, out: Tensor) -> int:
+def macs_adaptive_maxpool(_: _AdaptiveMaxPoolNd, inp: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.pooling._AdaptiveMaxPoolNd`"""
-
     # Approximate kernel_size using ratio of spatial shapes between input and output
     kernel_size = tuple(
         i_size // o_size if (i_size % o_size) == 0 else i_size - o_size * (i_size // o_size) + 1
@@ -147,9 +140,8 @@ def macs_adaptive_maxpool(module: _AdaptiveMaxPoolNd, inp: Tensor, out: Tensor)
     return out.numel() * (reduce(mul, kernel_size) - 1)
 
 
-def macs_adaptive_avgpool(module: _AdaptiveAvgPoolNd, inp: Tensor, out: Tensor) -> int:
+def macs_adaptive_avgpool(_: _AdaptiveAvgPoolNd, inp: Tensor, out: Tensor) -> int:
     """MACs estimation for `torch.nn.modules.pooling._AdaptiveAvgPoolNd`"""
-
     # Approximate kernel_size using ratio of spatial shapes between input and output
     kernel_size = tuple(
         i_size // o_size if (i_size % o_size) == 0 else i_size - o_size * (i_size // o_size) + 1
diff --git a/torchscan/modules/memory.py b/torchscan/modules/memory.py
index 2d73157..6fd6bbe 100644
--- a/torchscan/modules/memory.py
+++ b/torchscan/modules/memory.py
@@ -28,7 +28,6 @@ def module_dmas(module: Module, inp: Tensor, out: Tensor) -> int:
     Returns:
         int: number of DMAs
     """
-
     if isinstance(module, nn.Identity):
         return dmas_identity(module, inp, out)
     elif isinstance(module, nn.Flatten):
@@ -56,7 +55,7 @@ def module_dmas(module: Module, inp: Tensor, out: Tensor) -> int:
     elif isinstance(module, nn.Dropout):
         return dmas_dropout(module, inp, out)
     else:
-        warnings.warn(f"Module type not supported: {module.__class__.__name__}")
+        warnings.warn(f"Module type not supported: {module.__class__.__name__}", stacklevel=1)
         return 0
 
 
@@ -68,25 +67,21 @@ def num_params(module: Module) -> int:
     Returns:
         int: number of parameter elements
     """
-
     return sum(p.data.numel() for p in module.parameters())
 
 
-def dmas_identity(module: nn.Identity, inp: Tensor, out: Tensor) -> int:
+def dmas_identity(_: nn.Identity, inp: Tensor, __: Tensor) -> int:
     """DMAs estimation for `torch.nn.Identity`"""
-
     return inp.numel()
 
 
-def dmas_flatten(module: nn.Flatten, inp: Tensor, out: Tensor) -> int:
+def dmas_flatten(_: nn.Flatten, inp: Tensor, __: Tensor) -> int:
     """DMAs estimation for `torch.nn.Flatten`"""
-
     return 2 * inp.numel()
 
 
 def dmas_linear(module: nn.Linear, inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.Linear`"""
-
     input_dma = inp.numel()
     # Access weight and bias
     ops_dma = num_params(module)
@@ -97,7 +92,6 @@ def dmas_linear(module: nn.Linear, inp: Tensor, out: Tensor) -> int:
 
 def dmas_relu(module: Union[nn.ReLU, nn.ReLU6], inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.ReLU`"""
-
     input_dma = inp.numel()
     output_dma = 0 if module.inplace else out.numel()
 
@@ -106,7 +100,6 @@ def dmas_relu(module: Union[nn.ReLU, nn.ReLU6], inp: Tensor, out: Tensor) -> int
 
 def dmas_act_single_param(module: Union[nn.ELU, nn.LeakyReLU], inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for activations with single parameter"""
-
     input_dma = inp.numel()
     # Access alpha, slope or other
     ops_dma = 1
@@ -115,9 +108,8 @@ def dmas_act_single_param(module: Union[nn.ELU, nn.LeakyReLU], inp: Tensor, out:
     return input_dma + ops_dma + output_dma
 
 
-def dmas_sigmoid(module: nn.Sigmoid, inp: Tensor, out: Tensor) -> int:
+def dmas_sigmoid(_: nn.Sigmoid, inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.Sigmoid`"""
-
     # Access for both exp
     input_dma = inp.numel()
     output_dma = out.numel()
@@ -125,9 +117,8 @@ def dmas_sigmoid(module: nn.Sigmoid, inp: Tensor, out: Tensor) -> int:
     return input_dma + output_dma
 
 
-def dmas_tanh(module: nn.Tanh, inp: Tensor, out: Tensor) -> int:
+def dmas_tanh(_: nn.Tanh, inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.Tanh`"""
-
     # Access for both exp
     input_dma = inp.numel() * 2
     output_dma = out.numel()
@@ -137,7 +128,6 @@ def dmas_tanh(module: nn.Tanh, inp: Tensor, out: Tensor) -> int:
 
 def dmas_dropout(module: nn.Dropout, inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.Dropout`"""
-
     input_dma = inp.numel()
 
     # Access sampling probability
@@ -150,7 +140,6 @@ def dmas_dropout(module: nn.Dropout, inp: Tensor, out: Tensor) -> int:
 
 def dmas_convtransposend(module: _ConvTransposeNd, inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.modules.conv._ConvTransposeNd`"""
-
     # Padding (# cf. https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py#L496-L532)
     # Access stride, padding and kernel_size
     in_padding = len(module.kernel_size) * 4
@@ -162,9 +151,8 @@ def dmas_convtransposend(module: _ConvTransposeNd, inp: Tensor, out: Tensor) ->
     return in_padding + out_padding + conv_dmas
 
 
-def dmas_convnd(module: _ConvNd, inp: Tensor, out: Tensor) -> int:
+def dmas_convnd(module: _ConvNd, _: Tensor, out: Tensor) -> int:
     """DMAs estimation for `torch.nn.modules.conv._ConvNd`"""
-
     # Each output element required K ** 2 memory access of each input channel
     input_dma = module.in_channels * reduce(mul, module.kernel_size) * out.numel()
     # Correct with groups
@@ -202,7 +190,6 @@ def dmas_bn(module: _BatchNorm, inp: Tensor, out: Tensor) -> int:
 
 def dmas_pool(module: Union[_MaxPoolNd, _AvgPoolNd], inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for spatial pooling modules"""
-
     # Resolve kernel size and stride size (can be stored as a single integer or a tuple)
     if isinstance(module.kernel_size, tuple):
         kernel_size = module.kernel_size
@@ -217,9 +204,8 @@ def dmas_pool(module: Union[_MaxPoolNd, _AvgPoolNd], inp: Tensor, out: Tensor) -
     return input_dma + output_dma
 
 
-def dmas_adaptive_pool(module: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor) -> int:
+def dmas_adaptive_pool(_: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor) -> int:
     """DMAs estimation for adaptive spatial pooling modules"""
-
     # Approximate kernel_size using ratio of spatial shapes between input and output
     kernel_size = tuple(
         i_size // o_size if (i_size % o_size) == 0 else i_size - o_size * (i_size // o_size) + 1
diff --git a/torchscan/modules/receptive.py b/torchscan/modules/receptive.py
index 8d366f4..25611ae 100644
--- a/torchscan/modules/receptive.py
+++ b/torchscan/modules/receptive.py
@@ -52,30 +52,29 @@ def module_rf(module: Module, inp: Tensor, out: Tensor) -> Tuple[float, float, f
     elif isinstance(module, (_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd)):
         return rf_adaptive_poolnd(module, inp, out)
     else:
-        warnings.warn(f"Module type not supported: {module.__class__.__name__}")
+        warnings.warn(f"Module type not supported: {module.__class__.__name__}", stacklevel=1)
         return 1.0, 1.0, 0.0
 
 
-def rf_convtransposend(module: _ConvTransposeNd, intput: Tensor, out: Tensor) -> Tuple[float, float, float]:
+def rf_convtransposend(module: _ConvTransposeNd, _: Tensor, __: Tensor) -> Tuple[float, float, float]:
     k = module.kernel_size[0] if isinstance(module.kernel_size, tuple) else module.kernel_size
     s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
     return -k, 1.0 / s, 0.0
 
 
-def rf_aggregnd(module: Union[_ConvNd, _MaxPoolNd, _AvgPoolNd], inp: Tensor, out: Tensor) -> Tuple[float, float, float]:
+def rf_aggregnd(module: Union[_ConvNd, _MaxPoolNd, _AvgPoolNd], _: Tensor, __: Tensor) -> Tuple[float, float, float]:
     k = module.kernel_size[0] if isinstance(module.kernel_size, tuple) else module.kernel_size
     if hasattr(module, "dilation"):
         d = module.dilation[0] if isinstance(module.dilation, tuple) else module.dilation
-        k = d * (k - 1) + 1  # type: ignore[operator]
+        k = d * (k - 1) + 1
     s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
     p = module.padding[0] if isinstance(module.padding, tuple) else module.padding
     return k, s, p  # type: ignore[return-value]
 
 
 def rf_adaptive_poolnd(
-    module: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor
+    _: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor
 ) -> Tuple[int, int, float]:
-
     stride = math.ceil(inp.shape[-1] / out.shape[-1])
     kernel_size = stride
     padding = (inp.shape[-1] - kernel_size * stride) / 2
diff --git a/torchscan/process/memory.py b/torchscan/process/memory.py
index 1a5b67c..0065cb0 100644
--- a/torchscan/process/memory.py
+++ b/torchscan/process/memory.py
@@ -4,7 +4,7 @@
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.
 
 import re
-import subprocess
+import subprocess  # noqa S404
 import warnings
 
 import torch
@@ -20,10 +20,9 @@ def get_process_gpu_ram(pid: int) -> float:
     Returns:
         RAM usage in Megabytes
     """
-
     # PyTorch is not responsible for GPU usage
     if not torch.cuda.is_available():
-        warnings.warn("CUDA is unavailable to PyTorch.")
+        warnings.warn("CUDA is unavailable to PyTorch.", stacklevel=1)
         return 0.0
 
     # Query the running processes on GPUs
@@ -40,8 +39,8 @@ def get_process_gpu_ram(pid: int) -> float:
             ["nvidia-smi", "--query-gpu=memory.used", "--format=csv"], capture_output=True
         ).stdout.decode()
         return float(res.split("\n")[1].split()[0])
-    except Exception as e:
-        warnings.warn(f"raised: {e}. Parsing NVIDIA-SMI failed.")
+    except FileNotFoundError as e:
+        warnings.warn(f"raised: {e}. Parsing NVIDIA-SMI failed.", stacklevel=1)
 
     # Default to overall RAM usage for this process on the GPU
     ram_str = torch.cuda.list_gpu_processes().split("\n")
diff --git a/torchscan/utils.py b/torchscan/utils.py
index 0ed8461..7c093e9 100644
--- a/torchscan/utils.py
+++ b/torchscan/utils.py
@@ -3,6 +3,7 @@
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.
 
+from itertools import starmap
 from typing import Any, Dict, List, Optional, Tuple
 
 
@@ -15,7 +16,6 @@ def format_name(name: str, depth: int = 0) -> str:
     Returns:
         formatted string
     """
-
     if depth == 0:
         return name
     elif depth == 1:
@@ -36,7 +36,6 @@ def wrap_string(s: str, max_len: int, delimiter: str = ".", wrap: str = "[...]",
     Returns:
         wrapped string
     """
-
     if len(s) <= max_len or mode is None:
         return s
 
@@ -58,7 +57,6 @@ def unit_scale(val: float) -> Tuple[float, str]:
     Returns:
         tuple of rescaled value and unit
     """
-
     if val // 1e12 > 0:
         return val / 1e12, "T"
     elif val // 1e9 > 0:
@@ -71,7 +69,8 @@ def unit_scale(val: float) -> Tuple[float, str]:
         return val, ""
 
 
-def format_s(f_string, min_w: Optional[int] = None, max_w: Optional[int] = None) -> str:
+def format_s(f_string: str, min_w: Optional[int] = None, max_w: Optional[int] = None) -> str:
+    """Format number strings"""
     if isinstance(min_w, int):
         f_string = f"{f_string:<{min_w}}"
     if isinstance(max_w, int):
@@ -87,25 +86,25 @@ def format_line_str(
     receptive_field: bool = False,
     effective_rf_stats: bool = False,
 ) -> List[str]:
-
+    """Wrap all information into multiple lines"""
     if not isinstance(col_w, list):
         col_w = [None] * 7  # type: ignore[list-item]
 
     max_len = col_w[0] + 3 if isinstance(col_w[0], int) else 100
     line_str = [
-        format_s(wrap_string(format_name(layer["name"], layer["depth"]), max_len, mode=wrap_mode), col_w[0], col_w[0])
+        format_s(wrap_string(format_name(layer["name"], layer["depth"]), max_len, mode=wrap_mode), col_w[0], col_w[0]),
+        format_s(layer["type"], col_w[1], col_w[1]),
+        format_s(str(layer["output_shape"]), col_w[2], col_w[2]),
+        format_s(f"{layer['grad_params'] + layer['nograd_params'] + layer['num_buffers']:,}", col_w[3], col_w[3]),
     ]
-    line_str.append(format_s(layer["type"], col_w[1], col_w[1]))
-    line_str.append(format_s(str(layer["output_shape"]), col_w[2], col_w[2]))
-    line_str.append(
-        format_s(f"{layer['grad_params'] + layer['nograd_params'] + layer['num_buffers']:,}", col_w[3], col_w[3])
-    )
 
     if receptive_field:
         line_str.append(format_s(f"{layer['rf']:.0f}", col_w[4], col_w[4]))
         if effective_rf_stats:
-            line_str.append(format_s(f"{layer['s']:.0f}", col_w[5], col_w[5]))
-            line_str.append(format_s(f"{layer['p']:.0f}", col_w[6], col_w[6]))
+            line_str.extend((
+                format_s(f"{layer['s']:.0f}", col_w[5], col_w[5]),
+                format_s(f"{layer['p']:.0f}", col_w[6], col_w[6]),
+            ))
 
     return line_str
 
@@ -123,7 +122,6 @@ def format_info(
     Returns:
         formatted information
     """
-
     # Set margin between cols
     margin = 4
     # Dynamic col width
@@ -141,7 +139,7 @@ def format_info(
         ]
 
     # Truncate columns that are too long
-    col_w = [min(v, max_v) for v, max_v in zip(col_w, max_w)]
+    col_w = list(starmap(min, zip(col_w, max_w)))
 
     if not receptive_field:
         col_w = col_w[:4]
@@ -159,9 +157,11 @@ def format_info(
     margin_str = " " * margin
 
     # Header
-    info_str = [thin_line]
-    info_str.append(margin_str.join([f"{col_name:<{col_w}}" for col_name, col_w in zip(headers, col_w)]))
-    info_str.append(thick_line)
+    info_str = [
+        thin_line,
+        margin_str.join([f"{col_name:<{col_w}}" for col_name, col_w in zip(headers, col_w)]),
+        thick_line,
+    ]
 
     # Layers
     for layer in module_info["layers"]:
@@ -169,12 +169,13 @@ def format_info(
         info_str.append((" " * margin).join(line_str))
 
     # Parameter information
-    info_str.append(thick_line)
-
-    info_str.append(f"Trainable params: {module_info['overall']['grad_params']:,}")
-    info_str.append(f"Non-trainable params: {module_info['overall']['nograd_params']:,}")
     num_params = module_info["overall"]["grad_params"] + module_info["overall"]["nograd_params"]
-    info_str.append(f"Total params: {num_params:,}")
+    info_str.extend((
+        thick_line,
+        f"Trainable params: {module_info['overall']['grad_params']:,}",
+        f"Non-trainable params: {module_info['overall']['nograd_params']:,}",
+        f"Total params: {num_params:,}",
+    ))
 
     # Static RAM usage
     info_str.append(dot_line)
@@ -183,9 +184,11 @@ def format_info(
     param_size = (module_info["overall"]["param_size"] + module_info["overall"]["buffer_size"]) / 1024**2
     overhead = module_info["overheads"]["framework"]["fwd"] + module_info["overheads"]["cuda"]["fwd"]
 
-    info_str.append(f"Model size (params + buffers): {param_size:.2f} Mb")
-    info_str.append(f"Framework & CUDA overhead: {overhead:.2f} Mb")
-    info_str.append(f"Total RAM usage: {param_size + overhead:.2f} Mb")
+    info_str.extend((
+        f"Model size (params + buffers): {param_size:.2f} Mb",
+        f"Framework & CUDA overhead: {overhead:.2f} Mb",
+        f"Total RAM usage: {param_size + overhead:.2f} Mb",
+    ))
 
     # FLOPS information
     info_str.append(dot_line)
@@ -194,11 +197,12 @@ def format_info(
     macs, macs_units = unit_scale(sum(layer["macs"] for layer in module_info["layers"]))
     dmas, dmas_units = unit_scale(sum(layer["dmas"] for layer in module_info["layers"]))
 
-    info_str.append(f"Floating Point Operations on forward: {flops:.2f} {flops_units}FLOPs")
-    info_str.append(f"Multiply-Accumulations on forward: {macs:.2f} {macs_units}MACs")
-    info_str.append(f"Direct memory accesses on forward: {dmas:.2f} {dmas_units}DMAs")
-
-    info_str.append(thin_line)
+    info_str.extend((
+        f"Floating Point Operations on forward: {flops:.2f} {flops_units}FLOPs",
+        f"Multiply-Accumulations on forward: {macs:.2f} {macs_units}MACs",
+        f"Direct memory accesses on forward: {dmas:.2f} {dmas_units}DMAs",
+        thin_line,
+    ))
 
     return "\n".join(info_str)
 
@@ -212,7 +216,6 @@ def aggregate_info(info: Dict[str, Any], max_depth: int) -> Dict[str, Any]:
     Returns:
         edited dictionary information
     """
-
     if not any(layer["depth"] == max_depth for layer in info["layers"]):
         raise ValueError("The `max_depth` argument cannot be higher than module depth.")