From c94f6dcd110e94a16fa2ac0934e2f0677f4f2b0d Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 11:17:39 -0400
Subject: [PATCH 01/16] fix(ci): trigger on tinyml-tinyverse changes too

tinyml-tinyverse's own test suite doesn't exist yet at this point in
history (added by the next commit) -- this adds only the path trigger.
The step that actually runs those tests is added once they exist, later
in this branch, so no commit in this sequence ever references a directory
that isn't there yet.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .github/workflows/test-modelmaker.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/test-modelmaker.yml b/.github/workflows/test-modelmaker.yml
index 7233d670..0c277e35 100644
--- a/.github/workflows/test-modelmaker.yml
+++ b/.github/workflows/test-modelmaker.yml
@@ -5,11 +5,13 @@ on:
     branches: [platypus_dev_1.3, main]
     paths:
       - 'tinyml-modelmaker/**'
+      - 'tinyml-tinyverse/**'
       - '.github/workflows/test-modelmaker.yml'
   pull_request:
     branches: [platypus_dev_1.3, main]
     paths:
       - 'tinyml-modelmaker/**'
+      - 'tinyml-tinyverse/**'
   workflow_dispatch:  # manual trigger
 
 jobs:

From 6b8d3c53d047f1fe260278ff35b05b2952e0ebe0 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 19:39:07 -0400
Subject: [PATCH 02/16] fix: validate torch.compile with a warmup pass so
 failures actually fall back to eager mode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

torch.compile() is lazy — the previous try/except only wrapped the wrap-time
call, not the deferred compilation that happens on first forward. A failing
compile (e.g. unsupported Triton/ptxas for the GPU's compute capability)
would crash mid-training instead of falling back to eager mode.

compile_model_if_enabled now accepts an optional input_shape and, when
provided, runs one warmup forward pass through the compiled model before
training starts. A failure at either the wrap step or the warmup step falls
back to the original, uncompiled model. Backward compatible: omitting
input_shape preserves the old (unguarded) behavior.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 tinyml-tinyverse/tests/__init__.py            |  0
 .../tests/test_compile_warmup_fallback.py     | 94 +++++++++++++++++++
 .../references/common/train_base.py           | 31 +++++-
 3 files changed, 122 insertions(+), 3 deletions(-)
 create mode 100644 tinyml-tinyverse/tests/__init__.py
 create mode 100644 tinyml-tinyverse/tests/test_compile_warmup_fallback.py

diff --git a/tinyml-tinyverse/tests/__init__.py b/tinyml-tinyverse/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tinyml-tinyverse/tests/test_compile_warmup_fallback.py b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
new file mode 100644
index 00000000..68a4f678
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
@@ -0,0 +1,94 @@
+from unittest.mock import patch, MagicMock
+import torch
+import torch.nn as nn
+import pytest
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class _FakeArgs:
+    def __init__(self, compile_model):
+        self.compile_model = compile_model
+
+
+def _get_logger():
+    import logging
+    return logging.getLogger("test_compile_warmup_fallback")
+
+
+def test_compile_disabled_returns_original_model():
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=0)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    assert result is model
+
+
+def test_compile_success_with_warmup_returns_compiled_model():
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    # torch.compile wraps in an OptimizedModule; on CPU with a trivial model
+    # this should succeed genuinely (no mocking needed — real compile on a
+    # tiny CPU model is fast and should not fail).
+    assert result is not None
+    out = result(torch.rand(1, 4))
+    assert out.shape == (1, 2)
+
+
+def test_warmup_failure_falls_back_to_original_model():
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1)
+
+    class _BrokenCompiledModel(nn.Module):
+        def forward(self, x):
+            raise RuntimeError("simulated Inductor/Triton compile failure")
+
+    with patch('torch.compile', return_value=_BrokenCompiledModel()):
+        result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+
+    # Must fall back to the ORIGINAL model, not the broken compiled one.
+    assert result is model
+    out = result(torch.rand(1, 4))
+    assert out.shape == (1, 2)
+
+
+def test_wrap_time_failure_falls_back_to_original_model():
+    """torch.compile() itself raising (not just the warmup forward) is still caught."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1)
+
+    with patch('torch.compile', side_effect=RuntimeError("simulated wrap-time failure")):
+        result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+
+    assert result is model
+
+
+def test_no_input_shape_skips_warmup_but_still_compiles():
+    """Backward compatibility: callers that don't pass input_shape get the
+    old behavior (compile attempted, no warmup, no new fallback coverage)."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1)
+    result = compile_model_if_enabled(model, args, _get_logger())  # no input_shape
+    assert result is not None
+
+
+def test_warmup_restores_original_training_mode():
+    """The warmup pass must not leave the model stuck in eval mode."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    model.train()
+    args = _FakeArgs(compile_model=1)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    assert result.training is True
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index c6a5a087..5d8c09e8 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -658,17 +658,31 @@ def move_model_to_device(model, device, logger):
         sys.exit(1)
 
 
-def compile_model_if_enabled(model, args, logger):
+def compile_model_if_enabled(model, args, logger, input_shape=None):
     """
     Apply torch.compile to the model if --compile-model is enabled.
 
     torch.compile (PyTorch 2.0+) fuses operations into optimized kernels,
     which can significantly speed up training (15-30% on supported backends).
 
+    torch.compile() itself is lazy — it does not compile anything until the
+    first forward call. To catch compile failures (e.g. a Triton/ptxas
+    version that doesn't yet support the GPU's compute capability) before
+    training starts rather than mid-epoch, this function runs one warmup
+    forward pass through the compiled model when input_shape is provided.
+    A failure at either the wrap step or the warmup step falls back to the
+    original, uncompiled model.
+
     Args:
         model: The model to potentially compile
         args: Parsed arguments (uses args.compile_model)
         logger: Logger instance
+        input_shape: Shape (including batch dim) of a representative input
+            tensor, e.g. (1,) + dataset.X.shape[1:]. Used to run a warmup
+            forward pass that validates compilation actually works on this
+            hardware/toolchain. If None, no warmup is performed and a
+            compile failure will surface later, unguarded, on the training
+            loop's first real forward pass (legacy behavior).
 
     Returns:
         The (possibly compiled) model
@@ -684,10 +698,21 @@ def compile_model_if_enabled(model, args, logger):
         else:
             backend = 'aot_eager'
         logger.info(f"Compiling model with torch.compile (backend={backend})")
+        original_model = model
         try:
-            model = torch.compile(model, backend=backend)
+            compiled_model = torch.compile(model, backend=backend)
+            if input_shape is not None:
+                device = next(compiled_model.parameters()).device if len(list(compiled_model.parameters())) > 0 else torch.device('cpu')
+                dummy_input = torch.rand(size=input_shape, device=device)
+                was_training = compiled_model.training
+                compiled_model.eval()
+                with torch.no_grad():
+                    compiled_model(dummy_input)
+                compiled_model.train(was_training)
+            model = compiled_model
         except Exception as e:
-            logger.warning(f"torch.compile failed, falling back to eager mode: {e}")
+            logger.warning(f"torch.compile failed (or failed its warmup pass), falling back to eager mode: {e}")
+            model = original_model
     return model
 
 

From 6f892fdb8abae84a3040b833ab0354c8f1f9ba38 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 19:44:35 -0400
Subject: [PATCH 03/16] fix: restore training mode unconditionally after
 torch.compile warmup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

torch.compile's OptimizedModule wraps the original model by reference,
so calling .eval() on the compiled wrapper during warmup also flips the
ORIGINAL model's .training flag. If the warmup forward pass raised, the
restore step was skipped (exception jumped straight to the outer except),
so the fallback model was returned stuck in eval mode. Wrap the
eval/forward/restore sequence in try/finally so restoration always runs.

Also strengthens test_warmup_failure_falls_back_to_original_model: the
old mock returned an unrelated standalone module that didn't share
state/reference with the original model, so it couldn't reproduce this
bug. The new version wraps the original model as a real child submodule
(mirroring OptimizedModule._orig_mod) and asserts the original model's
.training flag is restored after a failed warmup — confirmed this fails
against the old buggy code and passes against the fix.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../tests/test_compile_warmup_fallback.py     | 37 ++++++++++++++++++-
 .../references/common/train_base.py           | 16 ++++++--
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_compile_warmup_fallback.py b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
index 68a4f678..54a26785 100644
--- a/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
+++ b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
@@ -45,19 +45,52 @@ def test_compile_success_with_warmup_returns_compiled_model():
 
 
 def test_warmup_failure_falls_back_to_original_model():
+    """Regression test for: torch.compile's OptimizedModule wraps the
+    original model BY REFERENCE (shares its parameters/state), so calling
+    .eval() on the compiled wrapper also flips the ORIGINAL model's
+    .training flag. If the warmup forward pass then raises, the original
+    model must still come back with its training mode restored — not
+    stuck in eval() because the restore step got skipped on the failure
+    path.
+
+    A plain mock that returns an unrelated standalone nn.Module does NOT
+    reproduce this, because it doesn't share state with the original model
+    the way OptimizedModule does. To actually catch a regression of this
+    bug, this test wraps the original model as a genuine child submodule
+    (registered via a real nn.Module attribute) so that calling
+    .train()/.eval() on the wrapper recurses into and mutates the original
+    model's .training flag too — exactly like OptimizedModule._orig_mod.
+    """
     from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
     model = _TinyModel()
+    model.train()  # known starting state
     args = _FakeArgs(compile_model=1)
 
-    class _BrokenCompiledModel(nn.Module):
+    class _SharedStateBrokenCompiledModel(nn.Module):
+        def __init__(self, wrapped):
+            super().__init__()
+            # Registering the original model as a real submodule means
+            # nn.Module.train()/.eval() on this wrapper recurses into it,
+            # mutating wrapped.training too — mirroring how
+            # torch._dynamo.OptimizedModule wraps the original model by
+            # reference via self._orig_mod.
+            self._orig_mod = wrapped
+
         def forward(self, x):
             raise RuntimeError("simulated Inductor/Triton compile failure")
 
-    with patch('torch.compile', return_value=_BrokenCompiledModel()):
+    broken = _SharedStateBrokenCompiledModel(model)
+
+    with patch('torch.compile', return_value=broken):
         result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
 
     # Must fall back to the ORIGINAL model, not the broken compiled one.
     assert result is model
+    # Must be restored to its pre-call training state, not left in eval()
+    # mode from the failed warmup pass — this is the exact defect a naive
+    # mock (that doesn't share state/reference with the original model)
+    # would fail to catch.
+    assert model.training is True
     out = result(torch.rand(1, 4))
     assert out.shape == (1, 2)
 
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index 5d8c09e8..92c40770 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -706,9 +706,19 @@ def compile_model_if_enabled(model, args, logger, input_shape=None):
                 dummy_input = torch.rand(size=input_shape, device=device)
                 was_training = compiled_model.training
                 compiled_model.eval()
-                with torch.no_grad():
-                    compiled_model(dummy_input)
-                compiled_model.train(was_training)
+                try:
+                    with torch.no_grad():
+                        compiled_model(dummy_input)
+                finally:
+                    # Must run on BOTH success and failure: torch.compile's
+                    # OptimizedModule wraps the original model BY REFERENCE
+                    # (shares its parameters/state), so compiled_model.eval()
+                    # above also flips the ORIGINAL model's .training flag.
+                    # If the warmup forward pass raises, control jumps to the
+                    # outer except block and returns original_model — if we
+                    # hadn't restored here first, that fallback model would
+                    # be returned stuck in eval mode.
+                    compiled_model.train(was_training)
             model = compiled_model
         except Exception as e:
             logger.warning(f"torch.compile failed (or failed its warmup pass), falling back to eager mode: {e}")

From 3fd2beddfc2a530ee5d57caafa11131faa7c5967 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 19:51:55 -0400
Subject: [PATCH 04/16] fix: pass input_shape to compile_model_if_enabled at
 all 4 call sites

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../references/timeseries_anomalydetection/train.py             | 2 +-
 .../references/timeseries_classification/train.py               | 2 +-
 .../tinyml_tinyverse/references/timeseries_forecasting/train.py | 2 +-
 .../tinyml_tinyverse/references/timeseries_regression/train.py  | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py
index ed64a8ed..78a5e089 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py
@@ -229,7 +229,7 @@ def main(gpu, args):
             return
 
         move_model_to_device(model, device, logger)
-        model = compile_model_if_enabled(model, args, logger)
+        model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:])
         criterion = nn.MSELoss()
 
         global _float_best_metric
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py
index e4233a4a..1d07f32d 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py
@@ -272,7 +272,7 @@ def main(gpu, args):
             return
 
         move_model_to_device(model, device, logger)
-        model = compile_model_if_enabled(model, args, logger)
+        model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:])
         criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)
 
         model, model_without_ddp, model_ema = setup_distributed_model(model, args, device)
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py
index d7172be1..a9b5c55b 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py
@@ -194,7 +194,7 @@ def main(gpu, args):
             return
 
         move_model_to_device(model, device, logger)
-        model = compile_model_if_enabled(model, args, logger)
+        model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:])
         criterion = nn.HuberLoss()
 
         global _float_best_metric
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py
index 98e49b02..4d08d0b1 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py
@@ -187,7 +187,7 @@ def main(gpu, args):
             args.output_int = False
 
         move_model_to_device(model, device, logger)
-        model = compile_model_if_enabled(model, args, logger)
+        model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:])
         global _float_best_metric
         sample_inputs = None
         sample_targets = None

From 57ff5866b57ef0e3c942973c66ce02b16a9a775f Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 20:06:46 -0400
Subject: [PATCH 05/16] docs: clarify compile warmup only validates one graph
 variant, not the training-mode graph

---
 .../references/common/train_base.py             | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index 92c40770..740ea0dc 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -673,16 +673,25 @@ def compile_model_if_enabled(model, args, logger, input_shape=None):
     A failure at either the wrap step or the warmup step falls back to the
     original, uncompiled model.
 
+    Note: the warmup runs in eval mode, no_grad, batch size 1, outside any
+    autocast context — it does not exercise the training-mode graph (with
+    gradients, the real batch size, and AMP autocast if enabled), which
+    dynamo compiles separately on first real use. This warmup catches
+    hardware/toolchain-level failures that occur regardless of graph
+    variant (e.g. ptxas rejecting the GPU architecture for any kernel), but
+    does not guarantee the training-mode graph will also compile cleanly.
+
     Args:
         model: The model to potentially compile
         args: Parsed arguments (uses args.compile_model)
         logger: Logger instance
         input_shape: Shape (including batch dim) of a representative input
             tensor, e.g. (1,) + dataset.X.shape[1:]. Used to run a warmup
-            forward pass that validates compilation actually works on this
-            hardware/toolchain. If None, no warmup is performed and a
-            compile failure will surface later, unguarded, on the training
-            loop's first real forward pass (legacy behavior).
+            forward pass that validates compilation works on this
+            hardware/toolchain for at least one graph variant. If None, no
+            warmup is performed and a compile failure will surface later,
+            unguarded, on the training loop's first real forward pass
+            (legacy behavior).
 
     Returns:
         The (possibly compiled) model

From d0f420a7697e19dc3ccfb351d60690c648cd6d52 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 21:39:06 -0400
Subject: [PATCH 06/16] fix: unwrap torch.compile wrapper before
 ONNX/TorchScript export
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

torch.compile() wraps a model in torch._dynamo.OptimizedModule, exposing
the original module at ._orig_mod. Neither torch.jit.trace (quantized
export) nor torch.onnx.export (float export) can trace a dynamo-optimized
module directly, so export_model() crashed with "Detected that you are
using FX to torch.jit.trace a dynamo-optimized function" whenever
compile_model=1 actually succeeded (the hardware-defaults feature
auto-enables this on CUDA). Unwrap via getattr(model, '_orig_mod', model)
once before the existing deepcopy — a no-op for uncompiled models.
---
 .../tests/test_export_model_unwrap.py         | 49 +++++++++++++++++++
 .../tinyml_tinyverse/common/utils/utils.py    |  6 +++
 2 files changed, 55 insertions(+)
 create mode 100644 tinyml-tinyverse/tests/test_export_model_unwrap.py

diff --git a/tinyml-tinyverse/tests/test_export_model_unwrap.py b/tinyml-tinyverse/tests/test_export_model_unwrap.py
new file mode 100644
index 00000000..daca0932
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_export_model_unwrap.py
@@ -0,0 +1,49 @@
+"""Regression test for: export_model() crashing on a torch.compile-wrapped
+model with `RuntimeError: Detected that you are using FX to torch.jit.trace
+a dynamo-optimized function. This is not supported at the moment.`
+
+Root cause: torch.compile() wraps a model in torch._dynamo.OptimizedModule,
+which exposes the original module at ._orig_mod. Neither torch.onnx.export
+nor torch.jit.trace (both used inside export_model, depending on whether
+quantization is enabled) can trace a dynamo-optimized module directly.
+"""
+import os
+import tempfile
+
+import torch
+import torch.nn as nn
+
+from tinyml_tinyverse.common.utils.utils import export_model
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+def test_export_model_handles_compiled_model_float_path():
+    """The non-quantized (else) branch: torch.onnx.export must not choke on
+    a compiled model."""
+    model = _TinyModel()
+    compiled_model = torch.compile(model, backend='aot_eager')
+    # Trigger real compilation so we're testing an actual OptimizedModule,
+    # not just the lazy uncalled wrapper.
+    compiled_model(torch.rand(1, 4))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        export_model(compiled_model, input_shape=(1, 4), output_dir=tmpdir, quantization=0)
+        assert os.path.exists(os.path.join(tmpdir, 'model.onnx'))
+
+
+def test_export_model_uncompiled_model_still_works():
+    """Backward compatibility: an ordinary, uncompiled model must still
+    export exactly as before (the getattr fallback is a no-op)."""
+    model = _TinyModel()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        export_model(model, input_shape=(1, 4), output_dir=tmpdir, quantization=0)
+        assert os.path.exists(os.path.join(tmpdir, 'model.onnx'))
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
index a8fb82fa..319ccca0 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
@@ -1626,6 +1626,12 @@ def export_model(model, input_shape, output_dir, opset_version=17, quantization=
     logger.debug(f"Quantization Mode: {quantization}, {type(quantization)}")
     logger.info(f'Exporting ONNX model from: {onnx_file}')
 
+    # torch.compile() wraps a model in torch._dynamo.OptimizedModule, exposing
+    # the original module at ._orig_mod. Neither torch.jit.trace (used below
+    # for quantized export) nor torch.onnx.export (used for float export) can
+    # trace a dynamo-optimized module directly, so unwrap first. This is a
+    # no-op for uncompiled models (getattr falls back to model itself).
+    model = getattr(model, '_orig_mod', model)
     model_copy = copy.deepcopy(model)
     model_copy = model_copy.to(device)
     if quantization:

From 4ff3e62fb6608c1b7c4f420525a2affec160bf8a Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 21:51:06 -0400
Subject: [PATCH 07/16] fix: recursively unwrap compiled submodules before
 export, not just top-level

Single-level getattr(model, '_orig_mod', model) missed the actual failure
shape: timeseries_classification always wraps the (already-compiled) model
inside NeuralNetworkWithPreprocess.model after compile_model_if_enabled runs,
so the compiled OptimizedModule ends up one level below the top-level model,
not at it. torch.onnx.export/torch.jit.trace still crashed on the nested
compiled submodule.

unwrap_compiled_submodules() walks the full submodule tree recursively,
replacing every torch.compile-wrapped module (top-level or nested) with its
original. No-op wherever nothing is compiled.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../tests/test_export_model_unwrap.py         | 54 ++++++++++++++++++-
 .../tinyml_tinyverse/common/utils/utils.py    | 27 ++++++++--
 2 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_export_model_unwrap.py b/tinyml-tinyverse/tests/test_export_model_unwrap.py
index daca0932..10792e01 100644
--- a/tinyml-tinyverse/tests/test_export_model_unwrap.py
+++ b/tinyml-tinyverse/tests/test_export_model_unwrap.py
@@ -13,7 +13,7 @@
 import torch
 import torch.nn as nn
 
-from tinyml_tinyverse.common.utils.utils import export_model
+from tinyml_tinyverse.common.utils.utils import export_model, unwrap_compiled_submodules
 
 
 class _TinyModel(nn.Module):
@@ -25,6 +25,19 @@ def forward(self, x):
         return self.linear(x)
 
 
+class _WrapperModel(nn.Module):
+    """Mirrors NeuralNetworkWithPreprocess: a compiled submodule nested one
+    level below the top-level model, not at the top level itself -- the
+    exact shape that broke a single-level `getattr(model, '_orig_mod', ...)`
+    unwrap in real timeseries_classification runs."""
+    def __init__(self, inner):
+        super().__init__()
+        self.model = inner
+
+    def forward(self, x):
+        return self.model(x)
+
+
 def test_export_model_handles_compiled_model_float_path():
     """The non-quantized (else) branch: torch.onnx.export must not choke on
     a compiled model."""
@@ -47,3 +60,42 @@ def test_export_model_uncompiled_model_still_works():
     with tempfile.TemporaryDirectory() as tmpdir:
         export_model(model, input_shape=(1, 4), output_dir=tmpdir, quantization=0)
         assert os.path.exists(os.path.join(tmpdir, 'model.onnx'))
+
+
+def test_export_model_handles_nested_compiled_submodule():
+    """Reproduces the real timeseries_classification failure: the compiled
+    model isn't the top-level module -- it's nested one level inside a
+    wrapper (NeuralNetworkWithPreprocess.model), because compile happens
+    before that wrapping is applied. A single-level unwrap misses this."""
+    inner = _TinyModel()
+    compiled_inner = torch.compile(inner, backend='aot_eager')
+    compiled_inner(torch.rand(1, 4))
+    wrapped = _WrapperModel(compiled_inner)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        export_model(wrapped, input_shape=(1, 4), output_dir=tmpdir, quantization=0)
+        assert os.path.exists(os.path.join(tmpdir, 'model.onnx'))
+
+
+def test_unwrap_compiled_submodules_top_level():
+    model = _TinyModel()
+    compiled = torch.compile(model, backend='aot_eager')
+    compiled(torch.rand(1, 4))
+    result = unwrap_compiled_submodules(compiled)
+    assert result is model
+
+
+def test_unwrap_compiled_submodules_nested():
+    inner = _TinyModel()
+    compiled_inner = torch.compile(inner, backend='aot_eager')
+    compiled_inner(torch.rand(1, 4))
+    wrapped = _WrapperModel(compiled_inner)
+    result = unwrap_compiled_submodules(wrapped)
+    assert result is wrapped
+    assert result.model is inner
+
+
+def test_unwrap_compiled_submodules_noop_when_nothing_compiled():
+    model = _TinyModel()
+    result = unwrap_compiled_submodules(model)
+    assert result is model
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
index 319ccca0..fc884ac7 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
@@ -1615,6 +1615,24 @@ def print_file_level_classification_summary(dataset, predicted, ground_truth,pha
     df = pd.DataFrame(results)
     logger_flcs.info(f'File-Level Classification Summary of {phase}:\n {tabulate(df, headers="keys", tablefmt="pretty")}')
 
+
+def unwrap_compiled_submodules(model):
+    """Recursively replace any torch.compile-wrapped module (including the
+    top-level model itself) with its original, uncompiled module.
+
+    torch.compile() wraps a module in torch._dynamo.OptimizedModule, which
+    exposes the original module at ._orig_mod. The compiled module isn't
+    always the top-level model passed in -- some reference scripts wrap an
+    already-compiled model inside another module afterward -- so this walks
+    the full submodule tree rather than only checking the top level. It is
+    a no-op wherever no torch.compile wrapping is present.
+    """
+    model = getattr(model, '_orig_mod', model)
+    for name, child in list(model.named_children()):
+        setattr(model, name, unwrap_compiled_submodules(child))
+    return model
+
+
 def export_model(model, input_shape, output_dir, opset_version=17, quantization=0,
                  example_input=None, generic_model=False, remove_hooks_for_jit=False):
     logger = getLogger("root.export_model")
@@ -1629,9 +1647,12 @@ def export_model(model, input_shape, output_dir, opset_version=17, quantization=
     # torch.compile() wraps a model in torch._dynamo.OptimizedModule, exposing
     # the original module at ._orig_mod. Neither torch.jit.trace (used below
     # for quantized export) nor torch.onnx.export (used for float export) can
-    # trace a dynamo-optimized module directly, so unwrap first. This is a
-    # no-op for uncompiled models (getattr falls back to model itself).
-    model = getattr(model, '_orig_mod', model)
+    # trace a dynamo-optimized module directly, so unwrap first. The compiled
+    # module isn't always at the top level -- e.g. timeseries_classification
+    # wraps it inside NeuralNetworkWithPreprocess.model after compiling -- so
+    # this walks the full submodule tree, not just the outermost model. It is
+    # a no-op wherever no torch.compile wrapping is present.
+    model = unwrap_compiled_submodules(model)
     model_copy = copy.deepcopy(model)
     model_copy = model_copy.to(device)
     if quantization:

From 7f9c0f7842ad3a1d5c2c1492a7e242f97570e4e2 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 23:09:41 -0400
Subject: [PATCH 08/16] =?UTF-8?q?fix:=20skip=20torch.compile=20when=20quan?=
 =?UTF-8?q?tization=20is=20enabled=20=E2=80=94=20FX=20tracing=20can't=20ha?=
 =?UTF-8?q?ndle=20a=20compiled=20model?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

prepare_qat_fx (FX symbolic tracing) crashes with "Detected that you are
using FX to symbolically trace a dynamo-optimized function" when the model
passed to it is torch.compile-wrapped (torch._dynamo.OptimizedModule).
compile_model_if_enabled now skips compiling entirely whenever
args.quantization is set, since compiling first would be discarded before
quantization prep runs anyway. Float training (quantization=0) is
unaffected.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 ...test_compile_skipped_under_quantization.py | 78 +++++++++++++++++++
 .../references/common/train_base.py           | 13 ++++
 2 files changed, 91 insertions(+)
 create mode 100644 tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py

diff --git a/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
new file mode 100644
index 00000000..fa4bfd29
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
@@ -0,0 +1,78 @@
+"""Regression test for: compile_model=1 combined with FX-based quantization
+(quantization=1 or 2) crashing at prepare_qat_fx time with
+`RuntimeError: Detected that you are using FX to symbolically trace a
+dynamo-optimized function. This is not supported at the moment.`
+
+Root cause: torch.compile() wraps the model in torch._dynamo.OptimizedModule
+before quantization_wrapped_model() runs prepare_qat_fx on it. FX symbolic
+tracing cannot trace a dynamo-optimized module at all. This is a different
+incompatibility than the ONNX/TorchScript export tracing issue (fixed
+separately in export_model()) -- unwrapping right before quantization would
+work mechanically but would mean compile provides no benefit for the rest
+of a quantized run's training, so this fix skips compiling entirely
+whenever quantization is enabled.
+"""
+import torch
+import torch.nn as nn
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class _FakeArgs:
+    def __init__(self, compile_model, quantization):
+        self.compile_model = compile_model
+        self.quantization = quantization
+
+
+def _get_logger():
+    import logging
+    return logging.getLogger("test_compile_skipped_under_quantization")
+
+
+def test_compile_skipped_when_quantization_enabled():
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1, quantization=2)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    # Must be the ORIGINAL model, not compiled -- an OptimizedModule here
+    # would go on to crash prepare_qat_fx's FX symbolic trace.
+    assert result is model
+    assert not hasattr(result, '_orig_mod')
+
+
+def test_compile_skipped_when_quantization_is_ptq_mode():
+    """quantization=1 (generic PTQ/QAT) hits the same FX-trace incompatibility
+    as quantization=2 (TINPU) -- both go through prepare_qat_fx."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1, quantization=1)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    assert result is model
+
+
+def test_compile_still_happens_for_float_training():
+    """Zero behavior change for the common case: quantization=0 (float
+    training) still compiles exactly as before."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1, quantization=0)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    assert result is not model  # genuinely compiled
+    out = result(torch.rand(1, 4))
+    assert out.shape == (1, 2)
+
+
+def test_compile_disabled_and_quantization_enabled_is_still_a_noop():
+    """When compile_model=0, quantization doesn't matter -- nothing changes."""
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=0, quantization=2)
+    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    assert result is model
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index 740ea0dc..f348a5cb 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -696,6 +696,19 @@ def compile_model_if_enabled(model, args, logger, input_shape=None):
     Returns:
         The (possibly compiled) model
     """
+    if getattr(args, 'quantization', 0):
+        # FX-based quantization (prepare_qat_fx) symbolically traces the model,
+        # which cannot trace a torch.compile-wrapped module at all -- a different
+        # incompatibility than the ONNX/TorchScript export tracing issue handled
+        # separately in export_model(). Skip compiling rather than compile and
+        # then immediately discard the benefit before quantization prep runs.
+        if getattr(args, 'compile_model', 0):
+            logger.info(
+                "compile_model is enabled but quantization is also enabled "
+                "(FX-based quantization cannot trace a compiled model) -- "
+                "skipping torch.compile for this run."
+            )
+        return model
     if getattr(args, 'compile_model', 0) and hasattr(torch, 'compile'):
         # Determine the best backend for the current device
         device_type = str(next(model.parameters()).device).split(':')[0] if len(list(model.parameters())) > 0 else 'cpu'

From 106e2f9a60abb458e808c39d929d0b6810491c69 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 23:12:30 -0400
Subject: [PATCH 09/16] test: strengthen skip-compile tests with torch.compile
 mock and INFO-log assertion

---
 ...test_compile_skipped_under_quantization.py | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
index fa4bfd29..40c94659 100644
--- a/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
+++ b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
@@ -12,6 +12,8 @@
 of a quantized run's training, so this fix skips compiling entirely
 whenever quantization is enabled.
 """
+from unittest.mock import patch
+
 import torch
 import torch.nn as nn
 
@@ -40,11 +42,25 @@ def test_compile_skipped_when_quantization_enabled():
     from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
     model = _TinyModel()
     args = _FakeArgs(compile_model=1, quantization=2)
-    result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
+    with patch('torch.compile') as mock_compile:
+        result = compile_model_if_enabled(model, args, _get_logger(), input_shape=(1, 4))
     # Must be the ORIGINAL model, not compiled -- an OptimizedModule here
     # would go on to crash prepare_qat_fx's FX symbolic trace.
     assert result is model
-    assert not hasattr(result, '_orig_mod')
+    # Proves the skip happens BEFORE compilation is attempted, not that
+    # compilation happened to fail or get discarded afterward.
+    mock_compile.assert_not_called()
+
+
+def test_compile_skip_logs_why(caplog):
+    import logging
+    from tinyml_tinyverse.references.common.train_base import compile_model_if_enabled
+    model = _TinyModel()
+    args = _FakeArgs(compile_model=1, quantization=2)
+    logger = _get_logger()
+    with caplog.at_level(logging.INFO, logger=logger.name):
+        compile_model_if_enabled(model, args, logger, input_shape=(1, 4))
+    assert any('quantization' in record.message.lower() for record in caplog.records)
 
 
 def test_compile_skipped_when_quantization_is_ptq_mode():

From e0d03c363ece315d06e47a68077d5a45e6f72158 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Tue, 28 Jul 2026 23:53:10 -0400
Subject: [PATCH 10/16] fix: strip torch.compile wrapper prefix from saved
 checkpoints

torch.compile() wraps a model in torch._dynamo.OptimizedModule. Since
setup_distributed_model only assigns model_without_ddp = model.module
under DDP, in the non-DDP case model_without_ddp IS the compiled
wrapper, so state_dict() emits every key prefixed _orig_mod. The
downstream float->quantization weight transfer (load_weights.py)
can't match those keys, silently falls back to strict=False, and
discards the entire float-trained result -- no exception, no fatal
error, the pipeline just quietly retrains from random init and
reports success.

save_checkpoint and resume_from_checkpoint now unwrap via a local
getattr(model, '_orig_mod', model) at each call site, so checkpoints
always carry uncompiled key names. This intentionally does not use
the existing unwrap_compiled_submodules() helper, which mutates the
model's submodule tree in place via setattr and would silently
un-compile the live training model.

Also adds weights_only=False to resume_from_checkpoint's torch.load
call. This is a separate, pre-existing latent bug this fix's own test
suite surfaced: checkpoint['args'] is an argparse.Namespace (every
train.py passes the full args object to save_checkpoint), which is
not on torch's default weights_only safe-globals list, so on torch
>=2.6 the --resume path already fails before reaching load_state_dict,
independent of the compile-wrapper issue. Matches the weights_only=False
convention already used at every other non-tensor-only torch.load()
call site in this codebase (load_weights.py, and the load_saved_model
paths in each task's train.py).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../test_checkpoint_unwrap_compiled_model.py  | 120 ++++++++++++++++++
 .../references/common/train_base.py           |  30 ++++-
 2 files changed, 145 insertions(+), 5 deletions(-)
 create mode 100644 tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
new file mode 100644
index 00000000..4f9c9059
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -0,0 +1,120 @@
+"""Regression test for: checkpoints saved from a torch.compile-wrapped model
+carrying _orig_mod.-prefixed keys, which the float->quantization weight
+transfer path (load_weights.py) cannot match -- it falls back to
+strict=False and silently discards the entire float-trained result. No
+exception is raised; the pipeline reports success while quietly retraining
+from random init.
+
+Root cause: setup_distributed_model sets model_without_ddp = model when not
+using DDP, so when compile_model_if_enabled succeeded upstream,
+model_without_ddp IS the torch._dynamo.OptimizedModule wrapper. state_dict()
+on it emits every key prefixed _orig_mod.
+"""
+import torch
+import torch.nn as nn
+
+from tinyml_tinyverse.references.common.train_base import save_checkpoint, resume_from_checkpoint
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class _FakeOptimizer:
+    def state_dict(self):
+        return {}
+
+    def load_state_dict(self, d):
+        pass
+
+
+class _FakeScheduler:
+    def state_dict(self):
+        return {}
+
+    def load_state_dict(self, d):
+        pass
+
+
+class _FakeArgs:
+    def __init__(self, resume):
+        self.resume = resume
+
+
+def test_save_checkpoint_strips_orig_mod_prefix_from_compiled_model():
+    model = _TinyModel()
+    compiled_model = torch.compile(model, backend='aot_eager')
+    compiled_model(torch.rand(1, 4))  # trigger real compilation
+
+    checkpoint = save_checkpoint(
+        compiled_model, _FakeOptimizer(), _FakeScheduler(), epoch=0, args=_FakeArgs(resume=None),
+    )
+    keys = list(checkpoint['model'].keys())
+    assert keys, "checkpoint has no keys at all"
+    assert not any(k.startswith('_orig_mod.') for k in keys), keys
+
+
+def test_save_checkpoint_uncompiled_model_unaffected():
+    """Backward compatibility: an ordinary, uncompiled model's checkpoint
+    keys are unchanged (no _orig_mod. prefix ever existed to strip)."""
+    model = _TinyModel()
+    checkpoint = save_checkpoint(
+        model, _FakeOptimizer(), _FakeScheduler(), epoch=0, args=_FakeArgs(resume=None),
+    )
+    assert set(checkpoint['model'].keys()) == set(model.state_dict().keys())
+
+
+def test_checkpoint_round_trips_into_a_fresh_uncompiled_model():
+    """The actual failure mode: save from a compiled model, load into the
+    (uncompiled) model used for the next training phase, and confirm the
+    real trained weights -- not random-init defaults -- are what land."""
+    source = _TinyModel()
+    with torch.no_grad():
+        source.linear.weight.fill_(3.14)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+
+    checkpoint = save_checkpoint(
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=0, args=_FakeArgs(resume=None),
+    )
+
+    target = _TinyModel()  # fresh, randomly initialized, NOT compiled
+    assert not torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 3.14))
+    target.load_state_dict(checkpoint['model'], strict=True)  # must not need strict=False
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 3.14))
+
+
+def test_resume_from_checkpoint_symmetric_with_compiled_model():
+    """resume_from_checkpoint (the --resume path) must be able to load a
+    checkpoint saved by save_checkpoint back into a still-compiled model,
+    using the same unwrap on both sides."""
+    import tempfile
+    import os
+
+    source = _TinyModel()
+    with torch.no_grad():
+        source.linear.weight.fill_(2.71)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+
+    checkpoint = save_checkpoint(
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=5, args=_FakeArgs(resume=None),
+    )
+
+    fresh = _TinyModel()
+    compiled_fresh = torch.compile(fresh, backend='aot_eager')
+    compiled_fresh(torch.rand(1, 4))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'checkpoint.pth')
+        torch.save(checkpoint, ckpt_path)
+        args = _FakeArgs(resume=ckpt_path)
+        args.device = 'cpu'
+        resume_from_checkpoint(compiled_fresh, _FakeOptimizer(), _FakeScheduler(), None, args)
+
+    assert torch.allclose(fresh.linear.weight, torch.full_like(fresh.linear.weight, 2.71))
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index f348a5cb..9a32e6a0 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -604,13 +604,23 @@ def resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema
         Updated args with start_epoch
     """
     if args.resume:
-        checkpoint = torch.load(args.resume, map_location=args.device)
-        model_without_ddp.load_state_dict(checkpoint['model'])
+        # weights_only=False: checkpoint['args'] is an argparse.Namespace (or,
+        # in tests, an equivalent stand-in), which is not on torch's default
+        # weights_only safe-globals list. Matches the convention already used
+        # for every other non-tensor-only torch.load() in this codebase (see
+        # load_weights.py and the per-task train.py load_saved_model paths).
+        checkpoint = torch.load(args.resume, map_location=args.device, weights_only=False)
+        # Symmetric with save_checkpoint's unwrap: checkpoints always carry
+        # uncompiled key names, so load into the unwrapped model regardless
+        # of whether it's currently wrapped by torch.compile.
+        resume_model = getattr(model_without_ddp, '_orig_mod', model_without_ddp)
+        resume_model.load_state_dict(checkpoint['model'])
         optimizer.load_state_dict(checkpoint['optimizer'])
         lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
         args.start_epoch = checkpoint['epoch'] + 1
         if model_ema:
-            model_ema.load_state_dict(checkpoint['model_ema'])
+            resume_ema = getattr(model_ema, '_orig_mod', model_ema)
+            resume_ema.load_state_dict(checkpoint['model_ema'])
     return args
 
 
@@ -811,15 +821,25 @@ def save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, mod
     Returns:
         dict: The checkpoint dictionary
     """
+    # torch.compile() wraps a model in torch._dynamo.OptimizedModule; when
+    # not using DDP, model_without_ddp IS that wrapper (setup_distributed_model
+    # only assigns model_without_ddp = model.module under DDP). state_dict() on
+    # a compiled model emits every key prefixed _orig_mod., which downstream
+    # weight-loading (load_weights.py, used for the float->quantization
+    # transfer) cannot match -- it silently falls back to strict=False and
+    # discards the entire result. Unwrap before saving so checkpoints always
+    # carry the original, uncompiled key names.
+    checkpoint_model = getattr(model_without_ddp, '_orig_mod', model_without_ddp)
     checkpoint = {
-        'model': model_without_ddp.state_dict(),
+        'model': checkpoint_model.state_dict(),
         'optimizer': optimizer.state_dict(),
         'lr_scheduler': lr_scheduler.state_dict(),
         'epoch': epoch,
         'args': args
     }
     if model_ema:
-        checkpoint['model_ema'] = model_ema.state_dict()
+        checkpoint_ema = getattr(model_ema, '_orig_mod', model_ema)
+        checkpoint['model_ema'] = checkpoint_ema.state_dict()
     if extra_data:
         checkpoint.update(extra_data)
     return checkpoint

From 1ed7effc8721151fc7fa748f018105de6360ced2 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 00:01:48 -0400
Subject: [PATCH 11/16] fix: strip _orig_mod prefix from EMA checkpoint keys
 too, symmetric save+load

ExponentialMovingAverage (AveragedModel) deep-copies its source model into
self.module, so when the source was already compiled, the OptimizedModule
wrapper ends up nested at model_ema.module._orig_mod -- not at
model_ema._orig_mod itself, where the prior fix's top-level getattr unwrap
couldn't reach it. Strip the substring from the resulting state_dict keys
instead (handles any nesting depth), and remap symmetrically on load by
matching against model_ema's own current key names stripped the same way.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../test_checkpoint_unwrap_compiled_model.py  | 59 +++++++++++++++++++
 .../references/common/train_base.py           | 29 +++++++--
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
index 4f9c9059..6faf7a9e 100644
--- a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -14,6 +14,7 @@
 import torch.nn as nn
 
 from tinyml_tinyverse.references.common.train_base import save_checkpoint, resume_from_checkpoint
+from tinyml_tinyverse.common.utils.utils import ExponentialMovingAverage
 
 
 class _TinyModel(nn.Module):
@@ -118,3 +119,61 @@ def test_resume_from_checkpoint_symmetric_with_compiled_model():
         resume_from_checkpoint(compiled_fresh, _FakeOptimizer(), _FakeScheduler(), None, args)
 
     assert torch.allclose(fresh.linear.weight, torch.full_like(fresh.linear.weight, 2.71))
+
+
+def test_save_checkpoint_strips_orig_mod_prefix_from_compiled_ema():
+    """ExponentialMovingAverage (AveragedModel) deep-copies its source model
+    into self.module -- so when the source was already compiled, the
+    OptimizedModule wrapper ends up nested at model_ema.module._orig_mod,
+    not at model_ema._orig_mod itself. A top-level unwrap can't reach it."""
+    model = _TinyModel()
+    compiled_model = torch.compile(model, backend='aot_eager')
+    compiled_model(torch.rand(1, 4))
+    model_ema = ExponentialMovingAverage(compiled_model, decay=0.99)
+
+    checkpoint = save_checkpoint(
+        compiled_model, _FakeOptimizer(), _FakeScheduler(), epoch=0,
+        args=_FakeArgs(resume=None), model_ema=model_ema,
+    )
+    keys = list(checkpoint['model_ema'].keys())
+    assert keys, "ema checkpoint has no keys at all"
+    assert not any('_orig_mod' in k for k in keys), keys
+
+
+def test_resume_from_checkpoint_symmetric_with_compiled_ema():
+    """The EMA analogue of test_resume_from_checkpoint_symmetric_with_compiled_model:
+    a checkpoint saved from a compiled model+EMA must load back into a fresh
+    compiled model+EMA, restoring the real EMA weight values."""
+    import tempfile
+    import os
+
+    source = _TinyModel()
+    with torch.no_grad():
+        source.linear.weight.fill_(1.5)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+    source_ema = ExponentialMovingAverage(compiled_source, decay=0.99)
+    with torch.no_grad():
+        for p in source_ema.module.parameters():
+            p.fill_(1.5)
+
+    checkpoint = save_checkpoint(
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
+        args=_FakeArgs(resume=None), model_ema=source_ema,
+    )
+
+    fresh = _TinyModel()
+    compiled_fresh = torch.compile(fresh, backend='aot_eager')
+    compiled_fresh(torch.rand(1, 4))
+    fresh_ema = ExponentialMovingAverage(compiled_fresh, decay=0.99)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'checkpoint.pth')
+        torch.save(checkpoint, ckpt_path)
+        args = _FakeArgs(resume=ckpt_path)
+        args.device = 'cpu'
+        resume_from_checkpoint(compiled_fresh, _FakeOptimizer(), _FakeScheduler(), fresh_ema, args)
+
+    ema_weight = dict(fresh_ema.module.named_parameters())['_orig_mod.linear.weight'] \
+        if hasattr(fresh_ema.module, '_orig_mod') else dict(fresh_ema.module.named_parameters())['linear.weight']
+    assert torch.allclose(ema_weight, torch.full_like(ema_weight, 1.5))
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index 9a32e6a0..798ad367 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -619,8 +619,21 @@ def resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema
         lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
         args.start_epoch = checkpoint['epoch'] + 1
         if model_ema:
-            resume_ema = getattr(model_ema, '_orig_mod', model_ema)
-            resume_ema.load_state_dict(checkpoint['model_ema'])
+            # Symmetric with the save-side key-substring strip: checkpoint
+            # keys never carry _orig_mod., but model_ema's OWN current keys
+            # might (if it's still compiled at resume time). Build a mapping
+            # from each of model_ema's current keys, stripped the same way,
+            # back to its real current key -- an identity mapping when
+            # uncompiled, a prefix-restoring one when compiled -- so the
+            # checkpoint's stripped keys land on whatever model_ema actually
+            # calls them right now.
+            live_keys_by_stripped_name = {
+                k.replace('_orig_mod.', ''): k for k in model_ema.state_dict().keys()
+            }
+            remapped_ema_state = {
+                live_keys_by_stripped_name.get(k, k): v for k, v in checkpoint['model_ema'].items()
+            }
+            model_ema.load_state_dict(remapped_ema_state)
     return args
 
 
@@ -838,8 +851,16 @@ def save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, mod
         'args': args
     }
     if model_ema:
-        checkpoint_ema = getattr(model_ema, '_orig_mod', model_ema)
-        checkpoint['model_ema'] = checkpoint_ema.state_dict()
+        # ExponentialMovingAverage (AveragedModel) deep-copies its source model
+        # into self.module -- so when the source was already compiled, the
+        # OptimizedModule wrapper ends up nested at model_ema.module._orig_mod,
+        # not at model_ema._orig_mod itself. A top-level getattr unwrap (as
+        # used for the main model above) can't reach it; strip the substring
+        # from the resulting state_dict keys instead, which handles the
+        # wrapper at whatever depth it's nested at.
+        checkpoint['model_ema'] = {
+            k.replace('_orig_mod.', ''): v for k, v in model_ema.state_dict().items()
+        }
     if extra_data:
         checkpoint.update(extra_data)
     return checkpoint

From 2268e22232f798287bdc718acb2cacb528131748 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 00:09:20 -0400
Subject: [PATCH 12/16] =?UTF-8?q?test:=20add=20asymmetric=20compiled-save/?=
 =?UTF-8?q?uncompiled-load=20EMA=20case=20=E2=80=94=20the=20real=20product?=
 =?UTF-8?q?ion=20shape?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../test_checkpoint_unwrap_compiled_model.py  | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
index 6faf7a9e..72f94113 100644
--- a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -177,3 +177,42 @@ def test_resume_from_checkpoint_symmetric_with_compiled_ema():
     ema_weight = dict(fresh_ema.module.named_parameters())['_orig_mod.linear.weight'] \
         if hasattr(fresh_ema.module, '_orig_mod') else dict(fresh_ema.module.named_parameters())['linear.weight']
     assert torch.allclose(ema_weight, torch.full_like(ema_weight, 1.5))
+
+
+def test_resume_from_checkpoint_ema_compiled_save_uncompiled_load():
+    """The actual production shape: EMA saved from a compiled source (keys
+    stripped at save time), then resumed into a run where EMA is NOT
+    compiled -- e.g. a later quantization phase that (per the separate
+    skip-compile-under-quantization fix) never compiles at all. The old
+    load-side code's top-level getattr couldn't reach EMA's nested wrapper,
+    so this asymmetric direction is the real discriminating case."""
+    import tempfile
+    import os
+
+    source = _TinyModel()
+    with torch.no_grad():
+        source.linear.weight.fill_(4.2)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+    source_ema = ExponentialMovingAverage(compiled_source, decay=0.99)
+    with torch.no_grad():
+        for p in source_ema.module.parameters():
+            p.fill_(4.2)
+
+    checkpoint = save_checkpoint(
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
+        args=_FakeArgs(resume=None), model_ema=source_ema,
+    )
+
+    fresh = _TinyModel()  # NOT compiled this time
+    fresh_ema = ExponentialMovingAverage(fresh, decay=0.99)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'checkpoint.pth')
+        torch.save(checkpoint, ckpt_path)
+        args = _FakeArgs(resume=ckpt_path)
+        args.device = 'cpu'
+        resume_from_checkpoint(fresh, _FakeOptimizer(), _FakeScheduler(), fresh_ema, args)
+
+    ema_weight = dict(fresh_ema.module.named_parameters())['linear.weight']
+    assert torch.allclose(ema_weight, torch.full_like(ema_weight, 4.2))

From aff21fda0263c9fdd91cc6321e05047d8869e7ea Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 11:18:18 -0400
Subject: [PATCH 13/16] fix(ci): run tinyml-tinyverse tests

tinyml-tinyverse/tests/ (added in this branch's earlier commits) was never
executed by CI. No extra install step needed -- tinyml-tinyverse is already
pip installed in the shared Install dependencies step.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .github/workflows/test-modelmaker.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/test-modelmaker.yml b/.github/workflows/test-modelmaker.yml
index 0c277e35..14789585 100644
--- a/.github/workflows/test-modelmaker.yml
+++ b/.github/workflows/test-modelmaker.yml
@@ -75,3 +75,7 @@ jobs:
       - name: Tier 2 — Pipeline Smoke Tests
         working-directory: tinyml-modelmaker
         run: python -m pytest tests/test_pipeline_smoke.py -v --tb=short
+
+      - name: tinyml-tinyverse Tests
+        working-directory: tinyml-tinyverse
+        run: python -m pytest tests/ -v --tb=short

From 0326b5044179f0855535cae1af86f7973c2d003c Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 01:56:33 -0400
Subject: [PATCH 14/16] fix: strip _orig_mod prefix in load_weights() as
 defense-in-depth

save_checkpoint/resume_from_checkpoint (train_base.py) already stop writing
_orig_mod.-prefixed keys going forward, but load_weights() -- the actual
consumer for the float->quantization --weights transfer
(timeseries_base.py) -- had no handling for the prefix at all, only for
'module.' (DDP). A checkpoint saved by older, unpatched code, or by any
other torch.compile-using caller not covered by the train_base.py fix,
would still hit the original silent strict=False fallback that discards
100% of weights with no exception.

Strip _orig_mod. from incoming checkpoint data before the existing
'module.' realignment runs (it's never meaningful to preserve or match
against, unlike 'module.'), and symmetrically remap onto the live model's
own current key names if that model is itself currently compiled --
mirroring the same symmetric approach already used for the EMA checkpoint
path.

Also adds a regression test that round-trips a checkpoint through the
actual load_weights() consumer (not just direct load_state_dict), closing
the gap noted in review: the fix was previously verified against this
consumer only via ad-hoc manual testing, not a permanent test.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../test_checkpoint_unwrap_compiled_model.py  | 25 +++++
 .../tests/test_load_weights_orig_mod.py       | 94 +++++++++++++++++++
 .../common/utils/load_weights.py              | 21 ++++-
 3 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 tinyml-tinyverse/tests/test_load_weights_orig_mod.py

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
index 72f94113..ea3682ac 100644
--- a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -15,6 +15,7 @@
 
 from tinyml_tinyverse.references.common.train_base import save_checkpoint, resume_from_checkpoint
 from tinyml_tinyverse.common.utils.utils import ExponentialMovingAverage
+from tinyml_tinyverse.common.utils.load_weights import load_weights
 
 
 class _TinyModel(nn.Module):
@@ -90,6 +91,30 @@ def test_checkpoint_round_trips_into_a_fresh_uncompiled_model():
     assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 3.14))
 
 
+def test_checkpoint_round_trips_through_the_real_load_weights_consumer():
+    """Same scenario as test_checkpoint_round_trips_into_a_fresh_uncompiled_model,
+    but through the ACTUAL production consumer of these checkpoints --
+    load_weights.load_weights(), used for the float->quantization --weights
+    transfer in timeseries_base.py -- rather than a raw load_state_dict call.
+    This is the function whose silent strict=False fallback originally
+    masked the bug (it printed a yellow warning and continued with 100% of
+    weights discarded, no exception, pipeline reported success)."""
+    source = _TinyModel()
+    with torch.no_grad():
+        source.linear.weight.fill_(6.28)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+
+    checkpoint = save_checkpoint(
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=0, args=_FakeArgs(resume=None),
+    )
+
+    target = _TinyModel()  # fresh, randomly initialized, NOT compiled
+    assert not torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 6.28))
+    load_weights(target, checkpoint['model'], state_dict_name=None)
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 6.28))
+
+
 def test_resume_from_checkpoint_symmetric_with_compiled_model():
     """resume_from_checkpoint (the --resume path) must be able to load a
     checkpoint saved by save_checkpoint back into a still-compiled model,
diff --git a/tinyml-tinyverse/tests/test_load_weights_orig_mod.py b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
new file mode 100644
index 00000000..fa3150b5
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
@@ -0,0 +1,94 @@
+"""Regression tests for load_weights() handling of torch.compile's
+_orig_mod. wrapper-artifact prefix, as a defense-in-depth complement to the
+save-side fix in train_base.py's save_checkpoint/resume_from_checkpoint
+(which now write checkpoints without the prefix in the first place).
+
+This covers the case load_weights() is the actual consumer for: the
+float->quantization --weights transfer
+(timeseries_base.py -> load_weights.load_weights(..., state_dict_name='model')),
+including checkpoints that might still carry _orig_mod. keys from an older,
+unpatched save, or from any other compile-using caller not covered by the
+train_base.py fix.
+"""
+import copy
+
+import torch
+import torch.nn as nn
+
+from tinyml_tinyverse.common.utils.load_weights import load_weights
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+def _filled(value):
+    model = _TinyModel()
+    with torch.no_grad():
+        model.linear.weight.fill_(value)
+        model.linear.bias.fill_(value)
+    return model
+
+
+def test_existing_module_prefix_alignment_still_works():
+    """Backward compatibility: the pre-existing 'module.' (DDP) prefix
+    realignment must be unaffected by the new _orig_mod. handling."""
+    source = _filled(7.0)
+    data = {f'module.{k}': v for k, v in source.state_dict().items()}
+
+    target = _TinyModel()
+    load_weights(target, data, state_dict_name=None)
+
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 7.0))
+
+
+def test_orig_mod_prefix_in_checkpoint_data_is_stripped():
+    """A checkpoint saved by older, unpatched code (or any other
+    torch.compile-using caller) with raw _orig_mod. keys must still load
+    correctly into a plain, uncompiled model."""
+    source = _filled(5.5)
+    data = {f'_orig_mod.{k}': v for k, v in source.state_dict().items()}
+
+    target = _TinyModel()
+    load_weights(target, data, state_dict_name=None)
+
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 5.5))
+
+
+def test_loading_into_a_currently_compiled_model():
+    """The live model being loaded INTO is itself currently torch.compile-
+    wrapped (its own state_dict keys carry _orig_mod.), and the checkpoint
+    data does not (the real shape produced by the already-fixed
+    save_checkpoint). Data must remap onto the model's actual current key
+    names, not fail to match them."""
+    source = _filled(2.25)
+    data = copy.deepcopy(source.state_dict())  # no prefix, as save_checkpoint now writes
+
+    target = _TinyModel()
+    compiled_target = torch.compile(target, backend='aot_eager')
+    compiled_target(torch.rand(1, 4))  # trigger real compilation
+
+    load_weights(compiled_target, data, state_dict_name=None)
+
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 2.25))
+
+
+def test_orig_mod_in_both_data_and_live_model():
+    """Both sides compiled: checkpoint data has stale _orig_mod. keys (e.g.
+    an old unpatched save) AND the live model being loaded into is itself
+    currently compiled. Must still transfer correctly."""
+    source = _filled(9.0)
+    data = {f'_orig_mod.{k}': v for k, v in source.state_dict().items()}
+
+    target = _TinyModel()
+    compiled_target = torch.compile(target, backend='aot_eager')
+    compiled_target(torch.rand(1, 4))
+
+    load_weights(compiled_target, data, state_dict_name=None)
+
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 9.0))
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
index aca21d19..84bad98e 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
@@ -84,13 +84,32 @@ def load_weights(model, pretrained, change_names_dict=None, keep_original_names=
     if load_error:
         # model did not load correctly. do any translation required.
         model_dict = model.state_dict()
+        model_keys = list(model_dict.keys())
+
+        # Strip torch.compile's _orig_mod. wrapper-artifact prefix from the
+        # incoming data first (e.g. a checkpoint saved by older code, before
+        # save_checkpoint started unwrapping it). Unlike 'module.' below,
+        # _orig_mod. is never meaningful to preserve or match against -- it's
+        # leftover from whichever side was torch.compile-wrapped when its
+        # state_dict was captured, not a real structural distinction.
+        data = {k.replace('_orig_mod.', ''): v for k, v in data.items()}
 
         # align the prefix 'module.' between model and data
-        model_prefix = 'module.' if 'module.' in list(model_dict.keys())[0] else ''
+        model_prefix = 'module.' if 'module.' in model_keys[0] else ''
         data_prefix = 'module.' if 'module.' in list(data.keys())[0] else ''
         data = {k.replace(data_prefix,model_prefix):v for k,v in data.items()} if data_prefix != '' \
             else {model_prefix+k:v for k,v in data.items()}
 
+        # If the live model being loaded INTO is itself currently
+        # torch.compile-wrapped, its real key names carry _orig_mod. too --
+        # remap data's (now wrapper-free) keys onto whatever the model
+        # actually calls them right now, mirroring the symmetric save/load
+        # unwrap already applied to the checkpoint main-model and EMA paths
+        # in train_base.py's save_checkpoint/resume_from_checkpoint.
+        if any('_orig_mod.' in k for k in model_keys):
+            live_keys_by_stripped_name = {k.replace('_orig_mod.', ''): k for k in model_keys}
+            data = {live_keys_by_stripped_name.get(k, k): v for k, v in data.items()}
+
         # change the name in pretrained data name to the given names
         if change_names_dict is not None:
             new_data = copy.deepcopy(data) if keep_original_names else {}

From 898bc4a453205ec118c8604bfc615d5f9baa0b8e Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 11:09:43 -0400
Subject: [PATCH 15/16] fix: resume_from_checkpoint unsafe deserialization +
 old-checkpoint compat

Two independent bugs in the same function, found by adversarial security
and code review:

1. Security (unsafe deserialization): the weights_only=False added to
   torch.load(args.resume, ...) disabled torch >=2.6's default unpickling
   restriction wholesale, rather than allowlisting only the one non-tensor
   type the checkpoint actually needs (argparse.Namespace, for
   checkpoint['args']). A crafted --resume checkpoint from an untrusted
   source could execute arbitrary code via a pickle __reduce__ payload.
   Fixed via torch.serialization.safe_globals([Namespace]).

2. Backward compatibility: the prior fix's unwrap only touched the live
   model, not the checkpoint data, so it silently assumed checkpoint['model']
   always has clean (unprefixed) keys. A checkpoint written before this
   session's fixes existed -- with _orig_mod. keys, from a compiled model --
   would now fail to load into an unwrapped model with a strict key
   mismatch, where it used to work by accident (both sides had matching
   prefixes). Fixed by extracting the same "strip from data, then remap
   onto the live model's actual current keys" pattern already used for the
   EMA branch into a shared _load_symmetric() helper, applied to both the
   main model and EMA uniformly.

Also strengthens existing checkpoint tests to assert on .bias and full
parameter sets, not just .weight -- a bug that only corrupted bias handling
would have passed undetected.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../test_checkpoint_unwrap_compiled_model.py  | 120 +++++++++++++++---
 .../tests/test_load_weights_orig_mod.py       |  17 ++-
 .../references/common/train_base.py           |  59 +++++----
 3 files changed, 148 insertions(+), 48 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
index ea3682ac..1ba5a691 100644
--- a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -48,6 +48,21 @@ def __init__(self, resume):
         self.resume = resume
 
 
+def _fill(model, value):
+    with torch.no_grad():
+        model.linear.weight.fill_(value)
+        model.linear.bias.fill_(value)
+    return model
+
+
+def _assert_all_params_equal(model, value):
+    """Check every parameter, not just .weight -- a bug that only corrupted
+    .bias handling would otherwise pass undetected."""
+    for name, param in model.named_parameters():
+        assert torch.allclose(param, torch.full_like(param, value)), \
+            f"{name} was not correctly transferred (expected all {value})"
+
+
 def test_save_checkpoint_strips_orig_mod_prefix_from_compiled_model():
     model = _TinyModel()
     compiled_model = torch.compile(model, backend='aot_eager')
@@ -75,9 +90,7 @@ def test_checkpoint_round_trips_into_a_fresh_uncompiled_model():
     """The actual failure mode: save from a compiled model, load into the
     (uncompiled) model used for the next training phase, and confirm the
     real trained weights -- not random-init defaults -- are what land."""
-    source = _TinyModel()
-    with torch.no_grad():
-        source.linear.weight.fill_(3.14)
+    source = _fill(_TinyModel(), 3.14)
     compiled_source = torch.compile(source, backend='aot_eager')
     compiled_source(torch.rand(1, 4))
 
@@ -88,7 +101,7 @@ def test_checkpoint_round_trips_into_a_fresh_uncompiled_model():
     target = _TinyModel()  # fresh, randomly initialized, NOT compiled
     assert not torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 3.14))
     target.load_state_dict(checkpoint['model'], strict=True)  # must not need strict=False
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 3.14))
+    _assert_all_params_equal(target, 3.14)
 
 
 def test_checkpoint_round_trips_through_the_real_load_weights_consumer():
@@ -99,9 +112,7 @@ def test_checkpoint_round_trips_through_the_real_load_weights_consumer():
     This is the function whose silent strict=False fallback originally
     masked the bug (it printed a yellow warning and continued with 100% of
     weights discarded, no exception, pipeline reported success)."""
-    source = _TinyModel()
-    with torch.no_grad():
-        source.linear.weight.fill_(6.28)
+    source = _fill(_TinyModel(), 6.28)
     compiled_source = torch.compile(source, backend='aot_eager')
     compiled_source(torch.rand(1, 4))
 
@@ -112,7 +123,7 @@ def test_checkpoint_round_trips_through_the_real_load_weights_consumer():
     target = _TinyModel()  # fresh, randomly initialized, NOT compiled
     assert not torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 6.28))
     load_weights(target, checkpoint['model'], state_dict_name=None)
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 6.28))
+    _assert_all_params_equal(target, 6.28)
 
 
 def test_resume_from_checkpoint_symmetric_with_compiled_model():
@@ -122,9 +133,7 @@ def test_resume_from_checkpoint_symmetric_with_compiled_model():
     import tempfile
     import os
 
-    source = _TinyModel()
-    with torch.no_grad():
-        source.linear.weight.fill_(2.71)
+    source = _fill(_TinyModel(), 2.71)
     compiled_source = torch.compile(source, backend='aot_eager')
     compiled_source(torch.rand(1, 4))
 
@@ -143,7 +152,7 @@ def test_resume_from_checkpoint_symmetric_with_compiled_model():
         args.device = 'cpu'
         resume_from_checkpoint(compiled_fresh, _FakeOptimizer(), _FakeScheduler(), None, args)
 
-    assert torch.allclose(fresh.linear.weight, torch.full_like(fresh.linear.weight, 2.71))
+    _assert_all_params_equal(fresh, 2.71)
 
 
 def test_save_checkpoint_strips_orig_mod_prefix_from_compiled_ema():
@@ -199,9 +208,9 @@ def test_resume_from_checkpoint_symmetric_with_compiled_ema():
         args.device = 'cpu'
         resume_from_checkpoint(compiled_fresh, _FakeOptimizer(), _FakeScheduler(), fresh_ema, args)
 
-    ema_weight = dict(fresh_ema.module.named_parameters())['_orig_mod.linear.weight'] \
-        if hasattr(fresh_ema.module, '_orig_mod') else dict(fresh_ema.module.named_parameters())['linear.weight']
-    assert torch.allclose(ema_weight, torch.full_like(ema_weight, 1.5))
+    for name, param in fresh_ema.module.named_parameters():
+        assert torch.allclose(param, torch.full_like(param, 1.5)), \
+            f"{name} was not correctly transferred (expected all 1.5)"
 
 
 def test_resume_from_checkpoint_ema_compiled_save_uncompiled_load():
@@ -239,5 +248,82 @@ def test_resume_from_checkpoint_ema_compiled_save_uncompiled_load():
         args.device = 'cpu'
         resume_from_checkpoint(fresh, _FakeOptimizer(), _FakeScheduler(), fresh_ema, args)
 
-    ema_weight = dict(fresh_ema.module.named_parameters())['linear.weight']
-    assert torch.allclose(ema_weight, torch.full_like(ema_weight, 4.2))
+    for name, param in fresh_ema.module.named_parameters():
+        assert torch.allclose(param, torch.full_like(param, 4.2)), \
+            f"{name} was not correctly transferred (expected all 4.2)"
+
+
+def test_resume_from_checkpoint_loads_old_format_checkpoint_with_orig_mod_keys():
+    """Backward compatibility: a checkpoint written before this fix existed
+    (or by any other torch.compile-using caller) has raw _orig_mod.-prefixed
+    keys in checkpoint['model'] -- resume_from_checkpoint must still load it
+    into a plain, uncompiled model, not raise a strict key-mismatch error."""
+    import tempfile
+    import os
+
+    old_style_checkpoint = {
+        'model': {'_orig_mod.linear.weight': torch.full((2, 4), 8.5),
+                   '_orig_mod.linear.bias': torch.full((2,), 8.5)},
+        'optimizer': {},
+        'lr_scheduler': {},
+        'epoch': 7,
+    }
+
+    target = _TinyModel()  # NOT compiled -- the real shape once a compile-hardening fix skips compile for this phase
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'checkpoint.pth')
+        torch.save(old_style_checkpoint, ckpt_path)
+        args = _FakeArgs(resume=ckpt_path)
+        args.device = 'cpu'
+        resume_from_checkpoint(target, _FakeOptimizer(), _FakeScheduler(), None, args)
+
+    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 8.5))
+    assert torch.allclose(target.linear.bias, torch.full_like(target.linear.bias, 8.5))
+    assert args.start_epoch == 8
+
+
+def test_resume_from_checkpoint_rejects_untrusted_pickle_payload():
+    """Security regression guard: resume_from_checkpoint must NOT accept
+    arbitrary pickled objects wholesale (i.e. must not silently be
+    weights_only=False in spirit). Only the one non-tensor type the
+    checkpoint legitimately needs (argparse.Namespace, for checkpoint['args'])
+    is allowlisted -- anything else in the pickle stream must still be
+    rejected by torch's weights_only safety check."""
+    import tempfile
+    import os
+
+    class _NotOnTheAllowlist:
+        """Standin for an attacker-controlled class with a malicious
+        __reduce__; the actual payload doesn't matter for this test, only
+        that torch.load refuses to construct instances of arbitrary,
+        non-allowlisted classes."""
+        def __reduce__(self):
+            return (self.__class__, ())
+
+    malicious_checkpoint = {
+        'model': _TinyModel().state_dict(),
+        'optimizer': {},
+        'lr_scheduler': {},
+        'epoch': 0,
+        'payload': _NotOnTheAllowlist(),
+    }
+
+    target = _TinyModel()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'checkpoint.pth')
+        torch.save(malicious_checkpoint, ckpt_path)
+        args = _FakeArgs(resume=ckpt_path)
+        args.device = 'cpu'
+        try:
+            resume_from_checkpoint(target, _FakeOptimizer(), _FakeScheduler(), None, args)
+            raised = False
+        except Exception:
+            raised = True
+
+    assert raised, (
+        "resume_from_checkpoint accepted a pickle payload containing a "
+        "non-allowlisted class -- the weights_only safety check is not "
+        "actually restricting unpickling."
+    )
diff --git a/tinyml-tinyverse/tests/test_load_weights_orig_mod.py b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
index fa3150b5..4276d917 100644
--- a/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
+++ b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
@@ -35,6 +35,15 @@ def _filled(value):
     return model
 
 
+def _assert_all_params_equal(model, value):
+    """Check every parameter, not just .weight -- a bug that only corrupted
+    .bias handling (a differently-shaped tensor, so a distinct code path
+    through key matching) would otherwise pass undetected."""
+    for name, param in model.named_parameters():
+        assert torch.allclose(param, torch.full_like(param, value)), \
+            f"{name} was not correctly transferred (expected all {value})"
+
+
 def test_existing_module_prefix_alignment_still_works():
     """Backward compatibility: the pre-existing 'module.' (DDP) prefix
     realignment must be unaffected by the new _orig_mod. handling."""
@@ -44,7 +53,7 @@ def test_existing_module_prefix_alignment_still_works():
     target = _TinyModel()
     load_weights(target, data, state_dict_name=None)
 
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 7.0))
+    _assert_all_params_equal(target, 7.0)
 
 
 def test_orig_mod_prefix_in_checkpoint_data_is_stripped():
@@ -57,7 +66,7 @@ def test_orig_mod_prefix_in_checkpoint_data_is_stripped():
     target = _TinyModel()
     load_weights(target, data, state_dict_name=None)
 
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 5.5))
+    _assert_all_params_equal(target, 5.5)
 
 
 def test_loading_into_a_currently_compiled_model():
@@ -75,7 +84,7 @@ def test_loading_into_a_currently_compiled_model():
 
     load_weights(compiled_target, data, state_dict_name=None)
 
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 2.25))
+    _assert_all_params_equal(target, 2.25)
 
 
 def test_orig_mod_in_both_data_and_live_model():
@@ -91,4 +100,4 @@ def test_orig_mod_in_both_data_and_live_model():
 
     load_weights(compiled_target, data, state_dict_name=None)
 
-    assert torch.allclose(target.linear.weight, torch.full_like(target.linear.weight, 9.0))
+    _assert_all_params_equal(target, 9.0)
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index 798ad367..76689f6b 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
@@ -72,7 +72,7 @@
 import platform
 import sys
 import timeit
-from argparse import ArgumentParser
+from argparse import ArgumentParser, Namespace
 from logging import getLogger
 
 import numpy as np
@@ -604,36 +604,41 @@ def resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema
         Updated args with start_epoch
     """
     if args.resume:
-        # weights_only=False: checkpoint['args'] is an argparse.Namespace (or,
-        # in tests, an equivalent stand-in), which is not on torch's default
-        # weights_only safe-globals list. Matches the convention already used
-        # for every other non-tensor-only torch.load() in this codebase (see
-        # load_weights.py and the per-task train.py load_saved_model paths).
-        checkpoint = torch.load(args.resume, map_location=args.device, weights_only=False)
-        # Symmetric with save_checkpoint's unwrap: checkpoints always carry
-        # uncompiled key names, so load into the unwrapped model regardless
-        # of whether it's currently wrapped by torch.compile.
-        resume_model = getattr(model_without_ddp, '_orig_mod', model_without_ddp)
-        resume_model.load_state_dict(checkpoint['model'])
+        # checkpoint['args'] is an argparse.Namespace (or, in tests, an
+        # equivalent stand-in). torch >=2.6 defaults torch.load to
+        # weights_only=True, which refuses to unpickle it -- but disabling
+        # the check entirely (weights_only=False) would also accept an
+        # attacker-crafted checkpoint's arbitrary __reduce__ payload as code
+        # to execute. Allowlist only the one non-tensor type this checkpoint
+        # actually needs instead.
+        with torch.serialization.safe_globals([Namespace]):
+            checkpoint = torch.load(args.resume, map_location=args.device)
+
+        # Checkpoints from any era may or may not carry a torch.compile
+        # _orig_mod. prefix on their keys (old saves did, when the model was
+        # compiled; save_checkpoint no longer writes it, but a checkpoint
+        # from before that fix, or from some other caller, still might).
+        # Strip it from the incoming data unconditionally -- it's never
+        # meaningful to preserve -- then remap onto whatever the live model
+        # actually calls its own keys right now, which carries _orig_mod.
+        # itself if it's currently compiled. Mirrors load_weights.py's
+        # identical handling for the same reason.
+        def _load_symmetric(live_module, checkpoint_state):
+            checkpoint_state = {k.replace('_orig_mod.', ''): v for k, v in checkpoint_state.items()}
+            live_keys = list(live_module.state_dict().keys())
+            if any('_orig_mod.' in k for k in live_keys):
+                live_keys_by_stripped_name = {k.replace('_orig_mod.', ''): k for k in live_keys}
+                checkpoint_state = {
+                    live_keys_by_stripped_name.get(k, k): v for k, v in checkpoint_state.items()
+                }
+            live_module.load_state_dict(checkpoint_state)
+
+        _load_symmetric(model_without_ddp, checkpoint['model'])
         optimizer.load_state_dict(checkpoint['optimizer'])
         lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
         args.start_epoch = checkpoint['epoch'] + 1
         if model_ema:
-            # Symmetric with the save-side key-substring strip: checkpoint
-            # keys never carry _orig_mod., but model_ema's OWN current keys
-            # might (if it's still compiled at resume time). Build a mapping
-            # from each of model_ema's current keys, stripped the same way,
-            # back to its real current key -- an identity mapping when
-            # uncompiled, a prefix-restoring one when compiled -- so the
-            # checkpoint's stripped keys land on whatever model_ema actually
-            # calls them right now.
-            live_keys_by_stripped_name = {
-                k.replace('_orig_mod.', ''): k for k in model_ema.state_dict().keys()
-            }
-            remapped_ema_state = {
-                live_keys_by_stripped_name.get(k, k): v for k, v in checkpoint['model_ema'].items()
-            }
-            model_ema.load_state_dict(remapped_ema_state)
+            _load_symmetric(model_ema, checkpoint['model_ema'])
     return args
 
 

From 0179c2d510b386a4803720e5523c329594674bf8 Mon Sep 17 00:00:00 2001
From: M Platypus <t5fkg8d44d@privaterelay.appleid.com>
Date: Wed, 29 Jul 2026 11:14:13 -0400
Subject: [PATCH 16/16] fix(tests): use real argparse.Namespace for checkpoint
 round-trip tests

Tests that do a genuine torch.save + resume_from_checkpoint (torch.load)
round trip were passing _FakeArgs (a lightweight test stand-in) as the
args parameter to save_checkpoint, which becomes checkpoint['args'] and
must therefore actually be unpickled by torch.load's weights_only safety
check. Only argparse.Namespace is allowlisted (matching real production
usage, where args always comes from ArgumentParser.parse_args()), so these
tests started failing once that check was genuinely enforced. Switched to
a real Namespace() for exactly the calls that get serialized to disk;
_FakeArgs remains fine for args objects that are only ever read directly
as a live function parameter (.resume/.device), never pickled.

Also moves the pickle-rejection test's "not on the allowlist" stand-in
class to module level -- pickle cannot serialize local/nested classes at
all, which made that test fail for an unrelated reason before it ever
reached the weights_only check it's meant to exercise.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../test_checkpoint_unwrap_compiled_model.py  | 37 +++++++++++++------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
index 1ba5a691..cc9e65bf 100644
--- a/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -10,6 +10,8 @@
 model_without_ddp IS the torch._dynamo.OptimizedModule wrapper. state_dict()
 on it emits every key prefixed _orig_mod.
 """
+from argparse import Namespace
+
 import torch
 import torch.nn as nn
 
@@ -48,6 +50,18 @@ def __init__(self, resume):
         self.resume = resume
 
 
+class _NotOnTheSafeGlobalsAllowlist:
+    """Standin for an attacker-controlled class with a malicious __reduce__;
+    the actual payload doesn't matter for the pickle-rejection test below,
+    only that torch.load refuses to construct instances of arbitrary,
+    non-allowlisted classes. Must be module-level, not defined inside the
+    test function -- pickle cannot serialize local/nested classes at all,
+    which would make the test fail for an unrelated reason before it ever
+    reached the weights_only check it's meant to exercise."""
+    def __reduce__(self):
+        return (self.__class__, ())
+
+
 def _fill(model, value):
     with torch.no_grad():
         model.linear.weight.fill_(value)
@@ -138,7 +152,10 @@ def test_resume_from_checkpoint_symmetric_with_compiled_model():
     compiled_source(torch.rand(1, 4))
 
     checkpoint = save_checkpoint(
-        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=5, args=_FakeArgs(resume=None),
+        # Namespace, not _FakeArgs: this ends up as checkpoint['args'], which
+        # torch.load's weights_only safety check must unpickle -- only
+        # argparse.Namespace is allowlisted, matching real production usage.
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=5, args=Namespace(),
     )
 
     fresh = _TinyModel()
@@ -192,8 +209,10 @@ def test_resume_from_checkpoint_symmetric_with_compiled_ema():
             p.fill_(1.5)
 
     checkpoint = save_checkpoint(
+        # Namespace, not _FakeArgs: see the identical note in
+        # test_resume_from_checkpoint_symmetric_with_compiled_model.
         compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
-        args=_FakeArgs(resume=None), model_ema=source_ema,
+        args=Namespace(), model_ema=source_ema,
     )
 
     fresh = _TinyModel()
@@ -234,8 +253,10 @@ def test_resume_from_checkpoint_ema_compiled_save_uncompiled_load():
             p.fill_(4.2)
 
     checkpoint = save_checkpoint(
+        # Namespace, not _FakeArgs: see the identical note in
+        # test_resume_from_checkpoint_symmetric_with_compiled_model.
         compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
-        args=_FakeArgs(resume=None), model_ema=source_ema,
+        args=Namespace(), model_ema=source_ema,
     )
 
     fresh = _TinyModel()  # NOT compiled this time
@@ -293,20 +314,12 @@ def test_resume_from_checkpoint_rejects_untrusted_pickle_payload():
     import tempfile
     import os
 
-    class _NotOnTheAllowlist:
-        """Standin for an attacker-controlled class with a malicious
-        __reduce__; the actual payload doesn't matter for this test, only
-        that torch.load refuses to construct instances of arbitrary,
-        non-allowlisted classes."""
-        def __reduce__(self):
-            return (self.__class__, ())
-
     malicious_checkpoint = {
         'model': _TinyModel().state_dict(),
         'optimizer': {},
         'lr_scheduler': {},
         'epoch': 0,
-        'payload': _NotOnTheAllowlist(),
+        'payload': _NotOnTheSafeGlobalsAllowlist(),
     }
 
     target = _TinyModel()