From bf9552d8532923e3e8f695676219129889a8fe21 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Thu, 4 Dec 2025 11:59:04 -0500
Subject: [PATCH 1/9] Add python-docx

---
 CHANGELOG.md                      |  5 ++
 install_miniforge.bash            |  3 +-
 tests/torch_example.py            | 13 +++-
 tests/torch_example_like_tflow.py | 98 +++++++++++++++++++++++++++++++
 4 files changed, 116 insertions(+), 3 deletions(-)
 create mode 100644 tests/torch_example_like_tflow.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1b2501c..2102e40 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,8 +11,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Update example Miniforge version to 25.11.0-0
+
 ### Added
 
+- Explicit Pip Packages
+  - python-docx
+
 ### Removed
 
 ### Deprecated
diff --git a/install_miniforge.bash b/install_miniforge.bash
index 15ae7b9..2744695 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -73,7 +73,7 @@ fi
 # -----
 
 EXAMPLE_PY_VERSION="3.13"
-EXAMPLE_MINI_VERSION="25.3.1-0"
+EXAMPLE_MINI_VERSION="25.11.0-0"
 EXAMPLE_INSTALLDIR="/opt/GEOSpyD"
 EXAMPLE_DATE=$(date +%F)
 usage() {
@@ -665,6 +665,7 @@ $PIP_INSTALL goes2go
 $PIP_INSTALL nco
 $PIP_INSTALL cdo
 $PIP_INSTALL ecmwf-opendata
+$PIP_INSTALL python-docx
 
 # some packages require a Fortran compiler. This sometimes isn't available
 # on macs (though usually is)
diff --git a/tests/torch_example.py b/tests/torch_example.py
index c8dbe14..1b45562 100644
--- a/tests/torch_example.py
+++ b/tests/torch_example.py
@@ -5,8 +5,17 @@
 
 
 dtype = torch.float
-device = torch.device("cpu")
-# device = torch.device("cuda:0") # Uncomment this to run on GPU
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+    print(f"GPU memory allocated: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
+    print(f"GPU memory cached: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")
+    print(f"CUDA device count: {torch.cuda.device_count()}")
+    print(f"Current CUDA device: {torch.cuda.current_device()}")
+    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+else:
+    device = torch.device("cpu")
+    print("CUDA not available, using CPU")
 
 # Create random input and output data
 x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
diff --git a/tests/torch_example_like_tflow.py b/tests/torch_example_like_tflow.py
new file mode 100644
index 0000000..278b9d6
--- /dev/null
+++ b/tests/torch_example_like_tflow.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+print("PyTorch version:", torch.__version__)
+print("CUDA available:", torch.cuda.is_available())
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load MNIST dataset
+transform = transforms.Compose([
+    transforms.ToTensor(),  # This automatically normalizes to [0, 1]
+])
+
+train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+
+# Define the model
+model = nn.Sequential(
+    nn.Flatten(),
+    nn.Linear(28 * 28, 128),
+    nn.ReLU(),
+    nn.Dropout(0.2),
+    nn.Linear(128, 10)
+).to(device)
+
+print(model)
+
+# Loss and optimizer
+criterion = nn.CrossEntropyLoss()  # Combines softmax and negative log likelihood
+optimizer = optim.Adam(model.parameters())
+
+# Training loop
+epochs = 5
+for epoch in range(epochs):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+
+        # Zero the gradients
+        optimizer.zero_grad()
+
+        # Forward pass
+        output = model(data)
+        loss = criterion(output, target)
+
+        # Backward pass and optimize
+        loss.backward()
+        optimizer.step()
+
+        # Statistics
+        running_loss += loss.item()
+        _, predicted = torch.max(output.data, 1)
+        total += target.size(0)
+        correct += (predicted == target).sum().item()
+
+    accuracy = 100 * correct / total
+    avg_loss = running_loss / len(train_loader)
+    print(f'Epoch {epoch + 1}/{epochs} - Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%')
+
+# Evaluation
+model.eval()
+correct = 0
+total = 0
+
+with torch.no_grad():
+    for data, target in test_loader:
+        data, target = data.to(device), target.to(device)
+        output = model(data)
+        _, predicted = torch.max(output.data, 1)
+        total += target.size(0)
+        correct += (predicted == target).sum().item()
+
+test_accuracy = 100 * correct / total
+print(f'\nTest Accuracy: {test_accuracy:.2f}%')
+
+# Get probabilities for first 5 test samples
+model.eval()
+with torch.no_grad():
+    test_data, _ = next(iter(test_loader))
+    test_data = test_data[:5].to(device)
+    logits = model(test_data)
+    probabilities = torch.softmax(logits, dim=1)
+    print("\nProbabilities for first 5 test samples:")
+    print(probabilities)

From c8a5e3c2bb1501b5cf63c099d83dd6c4bc67ff29 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Wed, 10 Dec 2025 10:28:30 -0500
Subject: [PATCH 2/9] Update example miniforge version to 25.11.0-1

---
 CHANGELOG.md           | 2 +-
 install_miniforge.bash | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2102e40..c5f1f2b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
-- Update example Miniforge version to 25.11.0-0
+- Update example Miniforge version to 25.11.0-1
 
 ### Added
 
diff --git a/install_miniforge.bash b/install_miniforge.bash
index 2744695..a99d338 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -73,7 +73,7 @@ fi
 # -----
 
 EXAMPLE_PY_VERSION="3.13"
-EXAMPLE_MINI_VERSION="25.11.0-0"
+EXAMPLE_MINI_VERSION="25.11.0-1"
 EXAMPLE_INSTALLDIR="/opt/GEOSpyD"
 EXAMPLE_DATE=$(date +%F)
 usage() {

From f055bdcd819e1082594cd873d926326f3c84211b Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Wed, 10 Dec 2025 11:57:54 -0500
Subject: [PATCH 3/9] Tensorflow not for Python 3.14, Basemap for 3.14 only in
 conda

---
 CHANGELOG.md           |  6 ++++++
 install_miniforge.bash | 13 +++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c5f1f2b..caec83d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,14 +12,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Update example Miniforge version to 25.11.0-1
+- Disable TensorFlow installation for Python 3.14 (see https://github.com/tensorflow/tensorflow/issues/102890)
 
 ### Added
 
+- Explicit Conda Packages
+  - basemap (latest version for Python 3.14 support only on conda-forge)
 - Explicit Pip Packages
   - python-docx
 
 ### Removed
 
+- Explicit Pip Packages
+  - basemap (latest version for Python 3.14 support only on conda-forge)
+
 ### Deprecated
 
 ## [25.3.1] - 2025-10-02
diff --git a/install_miniforge.bash b/install_miniforge.bash
index a99d338..b7b278a 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -605,6 +605,8 @@ $PACKAGE_INSTALL uxarray
 
 $PACKAGE_INSTALL rasterio contextily
 
+$PACKAGE_INSTALL basemap
+
 # Only install pythran on linux. On mac it brings in an old clang
 if [[ $MINIFORGE_ARCH == Linux ]]
 then
@@ -651,14 +653,21 @@ $PIP_INSTALL PyRTF3 pipenv pymp-pypi h5py
 $PIP_INSTALL pycircleci metpy siphon questionary xgrads
 $PIP_INSTALL ruamel.yaml
 $PIP_INSTALL xgboost
-$PIP_INSTALL tensorflow evidential-deep-learning silence_tensorflow
+
+# Tensorflow does not support Python 3.14 yet
+# https://github.com/tensorflow/tensorflow/issues/102890
+if [[ $PYTHON_VER_WITHOUT_DOT -ge 314 ]]
+then
+   echo "Skipping tensorflow installation as Python $PYTHON_VER is 3.14 or higher"
+else
+   $PIP_INSTALL tensorflow evidential-deep-learning silence_tensorflow
+fi
 $PIP_INSTALL torch
 $PIP_INSTALL yaplon
 $PIP_INSTALL lxml
 $PIP_INSTALL juliandate
 $PIP_INSTALL pybufrkit
 $PIP_INSTALL pyephem
-$PIP_INSTALL basemap
 $PIP_INSTALL redis
 $PIP_INSTALL Flask
 $PIP_INSTALL goes2go

From 4f8521ef2b23a8f597e727d93c6dc7243ea8ca3f Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Wed, 10 Dec 2025 11:58:33 -0500
Subject: [PATCH 4/9] Update example Python version to 3.14

---
 CHANGELOG.md           | 1 +
 install_miniforge.bash | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index caec83d..65367bb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Update example Python version to 3.14
 - Update example Miniforge version to 25.11.0-1
 - Disable TensorFlow installation for Python 3.14 (see https://github.com/tensorflow/tensorflow/issues/102890)
 
diff --git a/install_miniforge.bash b/install_miniforge.bash
index b7b278a..3e9dd58 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -72,7 +72,7 @@ fi
 # Usage
 # -----
 
-EXAMPLE_PY_VERSION="3.13"
+EXAMPLE_PY_VERSION="3.14"
 EXAMPLE_MINI_VERSION="25.11.0-1"
 EXAMPLE_INSTALLDIR="/opt/GEOSpyD"
 EXAMPLE_DATE=$(date +%F)

From f1fd3ff82db6d24933ce3a3332631f4a9d610792 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Wed, 11 Mar 2026 08:43:43 -0400
Subject: [PATCH 5/9] Fix ffnet build for Python 3.13+ and mamba 2.x
 compatibility

- Remove --no-use-pep517 flag (dropped in pip 23.1)
- Add --no-build-isolation for Python 3.13+ so numpy is visible to
  pip's build backend during the ffnet Fortran compilation
- Fix misleading echo message in ffnet install block
- Remove --show-channel-urls from mamba list calls (unsupported in mamba 2.x)
- Update example Miniforge version to 26.1.0-0
- Update CHANGELOG
---
 CHANGELOG.md           | 12 +++++++++++-
 install_miniforge.bash | 18 ++++++++++++------
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 65367bb..29818a9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,10 +9,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix `ffnet` installation for Python 3.13+ and modern pip (26.x):
+  - Remove `--no-use-pep517` flag, which was dropped in pip 23.1
+  - Add `--no-build-isolation` for Python 3.13+ so that `numpy` (required at
+    Fortran compile time) is visible to pip's build backend instead of being
+    hidden inside an isolated sandbox
+  - Simplify misleading echo message in the ffnet install block
+- Remove `--show-channel-urls` flag from `mamba list` calls; the flag is not
+  supported by mamba 2.x and caused the end-of-install package listing to be
+  skipped with a spurious warning
+
 ### Changed
 
 - Update example Python version to 3.14
-- Update example Miniforge version to 25.11.0-1
+- Update example Miniforge version to 26.1.0-0
 - Disable TensorFlow installation for Python 3.14 (see https://github.com/tensorflow/tensorflow/issues/102890)
 
 ### Added
diff --git a/install_miniforge.bash b/install_miniforge.bash
index 3e9dd58..a83af9e 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -73,7 +73,7 @@ fi
 # -----
 
 EXAMPLE_PY_VERSION="3.14"
-EXAMPLE_MINI_VERSION="25.11.0-1"
+EXAMPLE_MINI_VERSION="26.1.0-0"
 EXAMPLE_INSTALLDIR="/opt/GEOSpyD"
 EXAMPLE_DATE=$(date +%F)
 usage() {
@@ -680,7 +680,7 @@ $PIP_INSTALL python-docx
 # on macs (though usually is)
 if [[ $FORTRAN_AVAILABLE == TRUE ]]
 then
-   echo "We have a Fortran compiler and are Python 3.12 or older. Installing ffnet"
+   echo "We have a Fortran compiler. Installing ffnet"
    # we need to install ffnet from https://github.com/mrkwjc/ffnet.git
    # This is because the version in PyPI is not compatible with Python 3
    # and latest scipy
@@ -692,8 +692,14 @@ then
    if [[ $PYTHON_VER_WITHOUT_DOT -ge 313 ]]
    then
       $PIP_INSTALL setuptools wheel
-      # We also need a new flag for Python 3.13
-      EXTRA_PIP_FLAGS='--no-use-pep517'
+   fi
+   # For Python 3.13+, pip's isolated build environment does not inherit the
+   # conda env packages (e.g. numpy), which ffnet needs at build time. Passing
+   # --no-build-isolation tells pip to use the already-installed packages from
+   # the conda env instead of creating a fresh isolated sandbox.
+   if [[ $PYTHON_VER_WITHOUT_DOT -ge 313 ]]
+   then
+      EXTRA_PIP_FLAGS='--no-build-isolation'
    else
       EXTRA_PIP_FLAGS=''
    fi
@@ -765,8 +771,8 @@ $PIP_INSTALL prompt_toolkit
 # Use mamba to output list of packages installed
 # ----------------------------------------------
 cd $MINIFORGE_ENVDIR
-$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME --show-channel-urls --explicit > distribution_spec_file.txt
-$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME --show-channel-urls > mamba_list_packages.txt
+$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME --explicit > distribution_spec_file.txt
+$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME > mamba_list_packages.txt
 ./bin/pip freeze > pip_freeze_packages.txt
 
 # Restore User's .mambarc and .condarc using cleanup function

From 563698067e201fc9522ab18abf9820de07aca98f Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Wed, 11 Mar 2026 08:46:29 -0400
Subject: [PATCH 6/9] Update READMEs for Miniforge 26.1.0-0 and Python 3.14

- Update example versions to Python 3.14 and Miniforge 26.1.0-0
- Fix install path format (prefix/version/date, not prefix/version_pyX.Y/date)
- Add missing --ffnet-hack option to usage in README.md
- Fix NOTE 1 in README.md (path is no longer based on Python version)
- Update README.GMAO, README.NAS, README.NCCS to use install_miniforge.bash
  (replace stale references to install_miniconda.bash)
---
 README.GMAO |  2 +-
 README.NAS  |  2 +-
 README.NCCS |  2 +-
 README.md   | 20 +++++++++++---------
 4 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/README.GMAO b/README.GMAO
index 8d423a5..aac2dd4 100644
--- a/README.GMAO
+++ b/README.GMAO
@@ -1,5 +1,5 @@
 The command for installing on ford1 is:
 
 ```
-$ ./install_miniconda.bash --python_version 3.12 --miniconda_version 24.5.0-0 --prefix /ford1/share/gmao_SIteam/GEOSpyD
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /ford1/share/gmao_SIteam/GEOSpyD
 ```
diff --git a/README.NAS b/README.NAS
index 50b487c..322f483 100644
--- a/README.NAS
+++ b/README.NAS
@@ -1,5 +1,5 @@
 The command for installing on NAS systems is:
 
 ```
-$ ./install_miniconda.bash --python_version 3.12 --miniconda_version 24.5.0-0 --prefix /nobackup/gmao_SIteam/GEOSpyD
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /nobackup/gmao_SIteam/GEOSpyD
 ```
diff --git a/README.NCCS b/README.NCCS
index 3bd4992..0a1b3a9 100644
--- a/README.NCCS
+++ b/README.NCCS
@@ -1,5 +1,5 @@
 The command for installing on NCCS system is:
 
 ```
-$ ./install_miniconda.bash --python_version 3.12 --miniconda_version 24.5.0-0 --prefix /usr/local/other/python/GEOSpyD/
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /usr/local/other/python/GEOSpyD/
 ```
diff --git a/README.md b/README.md
index 5039432..7590f90 100644
--- a/README.md
+++ b/README.md
@@ -14,12 +14,12 @@ to prevent infection from the Anaconda `defaults` channel, the script at the end
 In order to use the install script, you can run:
 
 ```
-./install_miniforge.bash --python_version 3.12 --miniforge_version 24.5.0-0 --prefix /opt/GEOSpyD
+./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /opt/GEOSpyD
 ```
 
 will create an install at:
 ```
-/opt/GEOSpyD/24.5.0-0_py3.12/YYYY-MM-DD
+/opt/GEOSpyD/26.1.0-0/YYYY-MM-DD
 ```
 
 where YYYY-MM-DD is the date of the install. We use a date so that if
@@ -28,17 +28,19 @@ the stack is re-installed, the previous install is not overwritten.
 ## Usage
 
 ```
-Usage: ./install_miniforge.bash --python_version <python version> --miniforge_version <miniforge> --prefix <prefix> [--micromamba | --mamba] [--blas <blas>]
+Usage: ./install_miniforge.bash --python_version <python version> --miniforge_version <miniforge> --prefix <prefix>
+                   [--micromamba | --mamba] [--blas <blas>] [--ffnet-hack]
 
    Required arguments:
-      --python_version <python version> (e.g., 3.12)
-      --miniforge_version <miniforge_version version> (e.g., 24.5.0-0)
+      --python_version <python version> (e.g., 3.14)
+      --miniforge_version <miniforge_version version> (e.g., 26.1.0-0)
       --prefix <full path to installation directory> (e.g, /opt/GEOSpyD)
 
    Optional arguments:
       --blas <blas> (default: accelerate, options: mkl, openblas, accelerate, blis)
       --micromamba: Use micromamba installer (default)
       --mamba: Use mamba installer
+      --ffnet-hack: Install ffnet from fork (used on Bucy due to odd issue not finding gfortran)
       --help: Print this message
 
    By default we use the micromamba installer on both Linux and macOS
@@ -47,13 +49,12 @@ Usage: ./install_miniforge.bash --python_version <python version> --miniforge_ve
    NOTE 1: This script installs within /opt/GEOSpyD with a path based on:
 
         1. The Miniforge version
-        2. The Python version
-        3. The date of the installation
+        2. The date of the installation
 
-   For example: ./install_miniforge.bash --python_version 3.12 --miniforge_version 24.5.0-0 --prefix /opt/GEOSpyD
+   For example: ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /opt/GEOSpyD
 
    will create an install at:
-       /opt/GEOSpyD/24.5.0-0_py3.12/2024-08-29
+       /opt/GEOSpyD/26.1.0-0/2026-03-11
 
   NOTE 2: This script will create or substitute a .mambarc
   and .condarc file in the user's home directory.  If you
@@ -61,3 +62,4 @@ Usage: ./install_miniforge.bash --python_version <python version> --miniforge_ve
   restored after installation.  We do this to ensure that the
   installation uses conda-forge as the default channel.
 ```
+

From 4593e93575fca1ca5b05fdb116e358cb0482a259 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Fri, 3 Apr 2026 09:10:09 -0400
Subject: [PATCH 7/9] Update PyTorch examples for MPS, add torchvision, and
 prepare 26.1.1-3 release

---
 .gitignore                        |   2 +
 CHANGELOG.md                      |  19 ++-
 README.GMAO                       |   2 +-
 README.NAS                        |   2 +-
 README.NCCS                       |   2 +-
 README.md                         |  10 +-
 install_miniforge.bash            |   4 +-
 tests/torch_example.py            | 106 +++++++++--------
 tests/torch_example_like_tflow.py | 189 +++++++++++++++++-------------
 9 files changed, 198 insertions(+), 138 deletions(-)

diff --git a/.gitignore b/.gitignore
index 0400778..c809687 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,3 +11,5 @@ Miniconda2/
 xor*
 tmp-for-ffnet/
 *.png
+
+data/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29818a9..1464e6e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+### Changed
+
+### Deprecated
+
+### Removed
+
+### Fixed
+
+## [26.1.1-3] - 2026-04-03
+
 ### Fixed
 
 - Fix `ffnet` installation for Python 3.13+ and modern pip (26.x):
@@ -21,12 +33,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Update PyTorch examples to support Apple Metal Performance Shaders (MPS)
 - Update example Python version to 3.14
-- Update example Miniforge version to 26.1.0-0
+- Update example Miniforge version to 26.1.1-3
 - Disable TensorFlow installation for Python 3.14 (see https://github.com/tensorflow/tensorflow/issues/102890)
 
 ### Added
 
+- Add `torchvision` to the `pip install` list in `install_miniforge.bash`
+- Add CPU vs. Accelerated (CUDA/MPS) comparisons to PyTorch tests (`tests/torch_example.py` and `tests/torch_example_like_tflow.py`)
 - Explicit Conda Packages
   - basemap (latest version for Python 3.14 support only on conda-forge)
 - Explicit Pip Packages
@@ -37,8 +52,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Explicit Pip Packages
   - basemap (latest version for Python 3.14 support only on conda-forge)
 
-### Deprecated
-
 ## [25.3.1] - 2025-10-02
 
 ### Changed
diff --git a/README.GMAO b/README.GMAO
index aac2dd4..b725817 100644
--- a/README.GMAO
+++ b/README.GMAO
@@ -1,5 +1,5 @@
 The command for installing on ford1 is:
 
 ```
-$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /ford1/share/gmao_SIteam/GEOSpyD
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.1-3 --prefix /ford1/share/gmao_SIteam/GEOSpyD
 ```
diff --git a/README.NAS b/README.NAS
index 322f483..9ff4c0a 100644
--- a/README.NAS
+++ b/README.NAS
@@ -1,5 +1,5 @@
 The command for installing on NAS systems is:
 
 ```
-$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /nobackup/gmao_SIteam/GEOSpyD
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.1-3 --prefix /nobackup/gmao_SIteam/GEOSpyD
 ```
diff --git a/README.NCCS b/README.NCCS
index 0a1b3a9..0f56fbb 100644
--- a/README.NCCS
+++ b/README.NCCS
@@ -1,5 +1,5 @@
 The command for installing on NCCS system is:
 
 ```
-$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /usr/local/other/python/GEOSpyD/
+$ ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.1-3 --prefix /usr/local/other/python/GEOSpyD/
 ```
diff --git a/README.md b/README.md
index 7590f90..7c276f0 100644
--- a/README.md
+++ b/README.md
@@ -14,12 +14,12 @@ to prevent infection from the Anaconda `defaults` channel, the script at the end
 In order to use the install script, you can run:
 
 ```
-./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /opt/GEOSpyD
+./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.1-3 --prefix /opt/GEOSpyD
 ```
 
 will create an install at:
 ```
-/opt/GEOSpyD/26.1.0-0/YYYY-MM-DD
+/opt/GEOSpyD/26.1.1-3/YYYY-MM-DD
 ```
 
 where YYYY-MM-DD is the date of the install. We use a date so that if
@@ -33,7 +33,7 @@ Usage: ./install_miniforge.bash --python_version <python version> --miniforge_ve
 
    Required arguments:
       --python_version <python version> (e.g., 3.14)
-      --miniforge_version <miniforge_version version> (e.g., 26.1.0-0)
+      --miniforge_version <miniforge_version version> (e.g., 26.1.1-3)
       --prefix <full path to installation directory> (e.g, /opt/GEOSpyD)
 
    Optional arguments:
@@ -51,10 +51,10 @@ Usage: ./install_miniforge.bash --python_version <python version> --miniforge_ve
         1. The Miniforge version
         2. The date of the installation
 
-   For example: ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.0-0 --prefix /opt/GEOSpyD
+   For example: ./install_miniforge.bash --python_version 3.14 --miniforge_version 26.1.1-3 --prefix /opt/GEOSpyD
 
    will create an install at:
-       /opt/GEOSpyD/26.1.0-0/2026-03-11
+       /opt/GEOSpyD/26.1.1-3/2026-03-11
 
   NOTE 2: This script will create or substitute a .mambarc
   and .condarc file in the user's home directory.  If you
diff --git a/install_miniforge.bash b/install_miniforge.bash
index a83af9e..2b32efe 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -73,7 +73,7 @@ fi
 # -----
 
 EXAMPLE_PY_VERSION="3.14"
-EXAMPLE_MINI_VERSION="26.1.0-0"
+EXAMPLE_MINI_VERSION="26.1.1-3"
 EXAMPLE_INSTALLDIR="/opt/GEOSpyD"
 EXAMPLE_DATE=$(date +%F)
 usage() {
@@ -662,7 +662,7 @@ then
 else
    $PIP_INSTALL tensorflow evidential-deep-learning silence_tensorflow
 fi
-$PIP_INSTALL torch
+$PIP_INSTALL torch torchvision
 $PIP_INSTALL yaplon
 $PIP_INSTALL lxml
 $PIP_INSTALL juliandate
diff --git a/tests/torch_example.py b/tests/torch_example.py
index 1b45562..e2afda6 100644
--- a/tests/torch_example.py
+++ b/tests/torch_example.py
@@ -2,53 +2,67 @@
 
 import torch
 import math
+import time
 
 
-dtype = torch.float
+def run_polynomial_regression(device_name):
+    device = torch.device(device_name)
+    print(f"\n========== Running on {device} ==========")
+    dtype = torch.float
+
+    # Create random input and output data
+    x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
+    y = torch.sin(x)
+
+    # Randomly initialize weights
+    a = torch.randn((), device=device, dtype=dtype)
+    b = torch.randn((), device=device, dtype=dtype)
+    c = torch.randn((), device=device, dtype=dtype)
+    d = torch.randn((), device=device, dtype=dtype)
+
+    learning_rate = 1e-6
+
+    start_time = time.time()
+    for t in range(2000):
+        # Forward pass: compute predicted y
+        y_pred = a + b * x + c * x**2 + d * x**3
+
+        # Compute and print loss
+        loss = (y_pred - y).pow(2).sum().item()
+
+        # Backprop to compute gradients of a, b, c, d with respect to loss
+        grad_y_pred = 2.0 * (y_pred - y)
+        grad_a = grad_y_pred.sum()
+        grad_b = (grad_y_pred * x).sum()
+        grad_c = (grad_y_pred * x**2).sum()
+        grad_d = (grad_y_pred * x**3).sum()
+
+        # Update weights using gradient descent
+        a -= learning_rate * grad_a
+        b -= learning_rate * grad_b
+        c -= learning_rate * grad_c
+        d -= learning_rate * grad_d
+
+    end_time = time.time()
+
+    print(f"Time taken:  {end_time - start_time:.4f} seconds")
+    print(f"Final loss:  {loss:.4f}")
+    print(
+        f"Result eq:   y = {a.item():.4f} + {b.item():.4f} x + {c.item():.4f} x^2 + {d.item():.4f} x^3"
+    )
+
+
+print(f"PyTorch version: {torch.__version__}")
+
+# Always run CPU first
+run_polynomial_regression("cpu")
+
+# Determine and run accelerated device
 if torch.cuda.is_available():
-    device = torch.device("cuda:0")
-    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
-    print(f"GPU memory allocated: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
-    print(f"GPU memory cached: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")
-    print(f"CUDA device count: {torch.cuda.device_count()}")
-    print(f"Current CUDA device: {torch.cuda.current_device()}")
-    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+    print(f"\nFound CUDA: {torch.cuda.get_device_name(0)}")
+    run_polynomial_regression("cuda:0")
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    print("\nFound Apple Metal Performance Shaders (MPS)")
+    run_polynomial_regression("mps")
 else:
-    device = torch.device("cpu")
-    print("CUDA not available, using CPU")
-
-# Create random input and output data
-x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
-y = torch.sin(x)
-
-# Randomly initialize weights
-a = torch.randn((), device=device, dtype=dtype)
-b = torch.randn((), device=device, dtype=dtype)
-c = torch.randn((), device=device, dtype=dtype)
-d = torch.randn((), device=device, dtype=dtype)
-
-learning_rate = 1e-6
-for t in range(2000):
-    # Forward pass: compute predicted y
-    y_pred = a + b * x + c * x ** 2 + d * x ** 3
-
-    # Compute and print loss
-    loss = (y_pred - y).pow(2).sum().item()
-    if t % 100 == 99:
-        print(t, loss)
-
-    # Backprop to compute gradients of a, b, c, d with respect to loss
-    grad_y_pred = 2.0 * (y_pred - y)
-    grad_a = grad_y_pred.sum()
-    grad_b = (grad_y_pred * x).sum()
-    grad_c = (grad_y_pred * x ** 2).sum()
-    grad_d = (grad_y_pred * x ** 3).sum()
-
-    # Update weights using gradient descent
-    a -= learning_rate * grad_a
-    b -= learning_rate * grad_b
-    c -= learning_rate * grad_c
-    d -= learning_rate * grad_d
-
-
-print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')
+    print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")
diff --git a/tests/torch_example_like_tflow.py b/tests/torch_example_like_tflow.py
index 278b9d6..af9393e 100644
--- a/tests/torch_example_like_tflow.py
+++ b/tests/torch_example_like_tflow.py
@@ -5,94 +5,125 @@
 import torch.optim as optim
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
+import time
 
 print("PyTorch version:", torch.__version__)
 print("CUDA available:", torch.cuda.is_available())
-
-# Set device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+if hasattr(torch.backends, "mps"):
+    print("MPS available:", torch.backends.mps.is_available())
 
 # Load MNIST dataset
-transform = transforms.Compose([
-    transforms.ToTensor(),  # This automatically normalizes to [0, 1]
-])
-
-train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
-test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+transform = transforms.Compose(
+    [
+        transforms.ToTensor(),  # This automatically normalizes to [0, 1]
+    ]
+)
+
+print("\nDownloading/Loading MNIST dataset...")
+train_dataset = datasets.MNIST(
+    root="./data", train=True, download=True, transform=transform
+)
+test_dataset = datasets.MNIST(
+    root="./data", train=False, download=True, transform=transform
+)
 
 train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
 test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
 
-# Define the model
-model = nn.Sequential(
-    nn.Flatten(),
-    nn.Linear(28 * 28, 128),
-    nn.ReLU(),
-    nn.Dropout(0.2),
-    nn.Linear(128, 10)
-).to(device)
-
-print(model)
-
-# Loss and optimizer
-criterion = nn.CrossEntropyLoss()  # Combines softmax and negative log likelihood
-optimizer = optim.Adam(model.parameters())
-
-# Training loop
-epochs = 5
-for epoch in range(epochs):
-    model.train()
-    running_loss = 0.0
+
+def run_mnist(device_name):
+    device = torch.device(device_name)
+    print(f"\n========== Running on {device} ==========")
+
+    # Define the model
+    model = nn.Sequential(
+        nn.Flatten(),
+        nn.Linear(28 * 28, 128),
+        nn.ReLU(),
+        nn.Dropout(0.2),
+        nn.Linear(128, 10),
+    ).to(device)
+
+    # Loss and optimizer
+    criterion = nn.CrossEntropyLoss()  # Combines softmax and negative log likelihood
+    optimizer = optim.Adam(model.parameters())
+
+    # Training loop
+    epochs = 5
+    start_time = time.time()
+
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        for batch_idx, (data, target) in enumerate(train_loader):
+            data, target = data.to(device), target.to(device)
+
+            # Zero the gradients
+            optimizer.zero_grad()
+
+            # Forward pass
+            output = model(data)
+            loss = criterion(output, target)
+
+            # Backward pass and optimize
+            loss.backward()
+            optimizer.step()
+
+            # Statistics
+            running_loss += loss.item()
+            _, predicted = torch.max(output.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+        accuracy = 100 * correct / total
+        avg_loss = running_loss / len(train_loader)
+        print(
+            f"Epoch {epoch + 1}/{epochs} - Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%"
+        )
+
+    end_time = time.time()
+    print(f"-> Total Training Time: {end_time - start_time:.4f} seconds")
+
+    # Evaluation
+    model.eval()
     correct = 0
     total = 0
 
-    for batch_idx, (data, target) in enumerate(train_loader):
-        data, target = data.to(device), target.to(device)
-
-        # Zero the gradients
-        optimizer.zero_grad()
-
-        # Forward pass
-        output = model(data)
-        loss = criterion(output, target)
-
-        # Backward pass and optimize
-        loss.backward()
-        optimizer.step()
-
-        # Statistics
-        running_loss += loss.item()
-        _, predicted = torch.max(output.data, 1)
-        total += target.size(0)
-        correct += (predicted == target).sum().item()
-
-    accuracy = 100 * correct / total
-    avg_loss = running_loss / len(train_loader)
-    print(f'Epoch {epoch + 1}/{epochs} - Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%')
-
-# Evaluation
-model.eval()
-correct = 0
-total = 0
-
-with torch.no_grad():
-    for data, target in test_loader:
-        data, target = data.to(device), target.to(device)
-        output = model(data)
-        _, predicted = torch.max(output.data, 1)
-        total += target.size(0)
-        correct += (predicted == target).sum().item()
-
-test_accuracy = 100 * correct / total
-print(f'\nTest Accuracy: {test_accuracy:.2f}%')
-
-# Get probabilities for first 5 test samples
-model.eval()
-with torch.no_grad():
-    test_data, _ = next(iter(test_loader))
-    test_data = test_data[:5].to(device)
-    logits = model(test_data)
-    probabilities = torch.softmax(logits, dim=1)
-    print("\nProbabilities for first 5 test samples:")
-    print(probabilities)
+    eval_start = time.time()
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            _, predicted = torch.max(output.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    eval_end = time.time()
+    test_accuracy = 100 * correct / total
+    print(
+        f"-> Test Accuracy: {test_accuracy:.2f}% (Eval Time: {eval_end - eval_start:.4f} seconds)"
+    )
+
+    # Get probabilities for first 5 test samples
+    with torch.no_grad():
+        test_data, _ = next(iter(test_loader))
+        test_data = test_data[:5].to(device)
+        logits = model(test_data)
+        probabilities = torch.softmax(logits, dim=1)
+        print("\nProbabilities for first 5 test samples:")
+        print(probabilities)
+
+
+# Run CPU
+run_mnist("cpu")
+
+# Determine and run accelerated device
+if torch.cuda.is_available():
+    run_mnist("cuda:0")
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    run_mnist("mps")
+else:
+    print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")

From 23738db7f71cca47406cf081ab55806e243a259e Mon Sep 17 00:00:00 2001
From: Matt Thompson <matthew.thompson@nasa.gov>
Date: Fri, 3 Apr 2026 09:19:34 -0400
Subject: [PATCH 8/9] Update install_miniforge.bash

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
---
 install_miniforge.bash | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/install_miniforge.bash b/install_miniforge.bash
index 2b32efe..8df1bb4 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -771,8 +771,8 @@ $PIP_INSTALL prompt_toolkit
 # Use mamba to output list of packages installed
 # ----------------------------------------------
 cd $MINIFORGE_ENVDIR
-$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME --explicit > distribution_spec_file.txt
-$MINIFORGE_BINDIR/mamba list -n $MINIFORGE_ENVNAME > mamba_list_packages.txt
+"$MINIFORGE_BINDIR"/mamba list -n "$MINIFORGE_ENVNAME" --explicit > distribution_spec_file.txt
+"$MINIFORGE_BINDIR"/mamba list -n "$MINIFORGE_ENVNAME" > mamba_list_packages.txt
 ./bin/pip freeze > pip_freeze_packages.txt
 
 # Restore User's .mambarc and .condarc using cleanup function

From 50be101410c69ab63101685ad500028326fe3365 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Fri, 3 Apr 2026 09:23:54 -0400
Subject: [PATCH 9/9] Apply CodeRabbit review suggestions

---
 README.md                         |  2 +-
 install_miniforge.bash            |  2 +-
 tests/torch_example.py            | 31 +++++++-----
 tests/torch_example_like_tflow.py | 79 ++++++++++++++++---------------
 4 files changed, 62 insertions(+), 52 deletions(-)

diff --git a/README.md b/README.md
index 7c276f0..139e7d1 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ In order to use the install script, you can run:
 ```
 
 will create an install at:
-```
+```text
 /opt/GEOSpyD/26.1.1-3/YYYY-MM-DD
 ```
 
diff --git a/install_miniforge.bash b/install_miniforge.bash
index 8df1bb4..e447aff 100755
--- a/install_miniforge.bash
+++ b/install_miniforge.bash
@@ -680,7 +680,7 @@ $PIP_INSTALL python-docx
 # on macs (though usually is)
 if [[ $FORTRAN_AVAILABLE == TRUE ]]
 then
-   echo "We have a Fortran compiler. Installing ffnet"
+   echo "We have a Fortran compiler. Installing ffnet (Python $PYTHON_VER)"
    # we need to install ffnet from https://github.com/mrkwjc/ffnet.git
    # This is because the version in PyPI is not compatible with Python 3
    # and latest scipy
diff --git a/tests/torch_example.py b/tests/torch_example.py
index e2afda6..7f3ffe2 100644
--- a/tests/torch_example.py
+++ b/tests/torch_example.py
@@ -23,7 +23,7 @@ def run_polynomial_regression(device_name):
     learning_rate = 1e-6
 
     start_time = time.time()
-    for t in range(2000):
+    for _t in range(2000):
         # Forward pass: compute predicted y
         y_pred = a + b * x + c * x**2 + d * x**3
 
@@ -52,17 +52,22 @@ def run_polynomial_regression(device_name):
     )
 
 
-print(f"PyTorch version: {torch.__version__}")
+def main():
+    print(f"PyTorch version: {torch.__version__}")
 
-# Always run CPU first
-run_polynomial_regression("cpu")
+    # Always run CPU first
+    run_polynomial_regression("cpu")
 
-# Determine and run accelerated device
-if torch.cuda.is_available():
-    print(f"\nFound CUDA: {torch.cuda.get_device_name(0)}")
-    run_polynomial_regression("cuda:0")
-elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-    print("\nFound Apple Metal Performance Shaders (MPS)")
-    run_polynomial_regression("mps")
-else:
-    print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")
+    # Determine and run accelerated device
+    if torch.cuda.is_available():
+        print(f"\nFound CUDA: {torch.cuda.get_device_name(0)}")
+        run_polynomial_regression("cuda:0")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        print("\nFound Apple Metal Performance Shaders (MPS)")
+        run_polynomial_regression("mps")
+    else:
+        print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/torch_example_like_tflow.py b/tests/torch_example_like_tflow.py
index af9393e..f73bac3 100644
--- a/tests/torch_example_like_tflow.py
+++ b/tests/torch_example_like_tflow.py
@@ -7,31 +7,8 @@
 from torch.utils.data import DataLoader
 import time
 
-print("PyTorch version:", torch.__version__)
-print("CUDA available:", torch.cuda.is_available())
-if hasattr(torch.backends, "mps"):
-    print("MPS available:", torch.backends.mps.is_available())
-
-# Load MNIST dataset
-transform = transforms.Compose(
-    [
-        transforms.ToTensor(),  # This automatically normalizes to [0, 1]
-    ]
-)
-
-print("\nDownloading/Loading MNIST dataset...")
-train_dataset = datasets.MNIST(
-    root="./data", train=True, download=True, transform=transform
-)
-test_dataset = datasets.MNIST(
-    root="./data", train=False, download=True, transform=transform
-)
-
-train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
-
-
-def run_mnist(device_name):
+
+def run_mnist(device_name, train_loader, test_loader):
     device = torch.device(device_name)
     print(f"\n========== Running on {device} ==========")
 
@@ -58,7 +35,7 @@ def run_mnist(device_name):
         correct = 0
         total = 0
 
-        for batch_idx, (data, target) in enumerate(train_loader):
+        for data, target in train_loader:
             data, target = data.to(device), target.to(device)
 
             # Zero the gradients
@@ -74,7 +51,7 @@ def run_mnist(device_name):
 
             # Statistics
             running_loss += loss.item()
-            _, predicted = torch.max(output.data, 1)
+            predicted = output.argmax(dim=1)
             total += target.size(0)
             correct += (predicted == target).sum().item()
 
@@ -97,7 +74,7 @@ def run_mnist(device_name):
         for data, target in test_loader:
             data, target = data.to(device), target.to(device)
             output = model(data)
-            _, predicted = torch.max(output.data, 1)
+            predicted = output.argmax(dim=1)
             total += target.size(0)
             correct += (predicted == target).sum().item()
 
@@ -117,13 +94,41 @@ def run_mnist(device_name):
         print(probabilities)
 
 
-# Run CPU
-run_mnist("cpu")
+def main():
+    print("PyTorch version:", torch.__version__)
+    print("CUDA available:", torch.cuda.is_available())
+    if hasattr(torch.backends, "mps"):
+        print("MPS available:", torch.backends.mps.is_available())
+
+    # Load MNIST dataset
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),  # This automatically normalizes to [0, 1]
+        ]
+    )
+
+    print("\nDownloading/Loading MNIST dataset...")
+    train_dataset = datasets.MNIST(
+        root="./data", train=True, download=True, transform=transform
+    )
+    test_dataset = datasets.MNIST(
+        root="./data", train=False, download=True, transform=transform
+    )
+
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+
+    # Run CPU
+    run_mnist("cpu", train_loader, test_loader)
+
+    # Determine and run accelerated device
+    if torch.cuda.is_available():
+        run_mnist("cuda:0", train_loader, test_loader)
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        run_mnist("mps", train_loader, test_loader)
+    else:
+        print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")
+
 
-# Determine and run accelerated device
-if torch.cuda.is_available():
-    run_mnist("cuda:0")
-elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-    run_mnist("mps")
-else:
-    print("\nNo accelerated device (CUDA/MPS) found. Skipping accelerated run.")
+if __name__ == "__main__":
+    main()