diff --git a/tinyml-modelmaker/tests/test_hardware_defaults.py b/tinyml-modelmaker/tests/test_hardware_defaults.py
new file mode 100644
index 00000000..84c25f88
--- /dev/null
+++ b/tinyml-modelmaker/tests/test_hardware_defaults.py
@@ -0,0 +1,54 @@
+from unittest.mock import patch
+import pytest
+
+
+def _make_params():
+    """Build a minimal ConfigDict with the two training flags at their defaults."""
+    from tinyml_modelmaker.utils.config_dict import ConfigDict
+    return ConfigDict(dict(training=dict(compile_model=0, native_amp=False)))
+
+
+def test_auto_enables_both_on_cuda_with_no_explicit_keys():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, set())
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_respects_explicit_native_amp_false():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, {'native_amp'})
+    assert params.training.compile_model == 1   # still auto-enabled
+    assert params.training.native_amp is False  # explicit choice respected
+
+
+def test_respects_explicit_compile_model_zero():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, {'compile_model'})
+    assert params.training.compile_model == 0   # explicit choice respected
+    assert params.training.native_amp is True   # still auto-enabled
+
+
+def test_no_change_without_cuda():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=False):
+        apply_hardware_defaults(params, set())
+    assert params.training.compile_model == 0
+    assert params.training.native_amp is False
+
+
+def test_safe_when_params_lacks_flags():
+    """hasattr guards: calling on a params dict without compile_model/native_amp must not raise."""
+    from tinyml_modelmaker.utils.config_dict import ConfigDict
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = ConfigDict(dict(training=dict(batch_size=32)))
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, set())  # must not raise
+    assert params.training.batch_size == 32
diff --git a/tinyml-modelmaker/tests/test_hardware_defaults_integration.py b/tinyml-modelmaker/tests/test_hardware_defaults_integration.py
new file mode 100644
index 00000000..8323725c
--- /dev/null
+++ b/tinyml-modelmaker/tests/test_hardware_defaults_integration.py
@@ -0,0 +1,69 @@
+from unittest.mock import patch
+import pytest
+
+
+def test_init_params_auto_enables_on_cuda():
+    """init_params with no training overrides auto-enables both flags on CUDA."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(common=dict(task_category='timeseries_classification'))
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params(user_config)
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_init_params_respects_native_amp_false_override():
+    """Explicit native_amp: false in user config is not overridden.
+
+    Mirrors the real production call pattern in run_tinyml_modelmaker.py:
+    ModelRunner.init_params() is called with zero arguments (so hardware
+    defaults auto-enable via apply_hardware_defaults with an empty
+    explicitly_set), and the user's config is merged in afterward via
+    ConfigDict.update() -- a deep merge -- not passed into the ConfigDict
+    constructor. (The constructor's *args merge path has a separate,
+    pre-existing shallow-merge bug that this test must avoid triggering.)
+    """
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(
+        common=dict(task_category='timeseries_classification'),
+        training=dict(native_amp=False),
+    )
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params()
+    params.update(user_config)
+    assert params.training.native_amp is False
+    assert params.training.compile_model == 1  # still auto-enabled
+
+
+def test_init_params_no_change_without_cuda():
+    """Without CUDA, both flags stay at defaults."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(common=dict(task_category='timeseries_classification'))
+    with patch('torch.cuda.is_available', return_value=False):
+        params = init_params(user_config)
+    assert params.training.compile_model == 0
+    assert params.training.native_amp is False
+
+
+def test_init_params_does_not_crash_on_non_dict_first_arg():
+    """ConfigDict documents a YAML path string (or None) as valid first
+    positional input, not just a dict. init_params's own explicitly-set-key
+    detection must not assume args[0] is dict-like and crash instead."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    with patch('torch.cuda.is_available', return_value=True):
+        # A bogus path is fine here -- ConfigDict's *args merge loop only
+        # merges dict/ConfigDict values and silently ignores strings, so
+        # this never actually attempts to read the file. The point of this
+        # test is solely that passing a string doesn't crash init_params's
+        # own explicitly-set-key detection before ConfigDict is even built.
+        params = init_params('/nonexistent/path/to/config.yaml')
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_init_params_does_not_crash_on_none_first_arg():
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params(None)
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/params.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/params.py
index dc891a67..c468eb1e 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/params.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/params.py
@@ -35,6 +35,7 @@
 
 from ... import utils
 from . import constants
+from ...utils.hardware_defaults import apply_hardware_defaults
 
 
 def init_params(*args, **kwargs):
@@ -222,5 +223,11 @@ def init_params(*args, **kwargs):
         ),
     )
 
+    # args[0] is usually a user config dict, but ConfigDict itself also
+    # accepts a YAML path string or None as a first positional argument --
+    # only inspect it as a mapping when it actually is one.
+    user_training_keys = set(args[0].get('training', {}).keys()) \
+        if args and isinstance(args[0], dict) else set()
     params = utils.ConfigDict(default_params, *args, **kwargs)
+    apply_hardware_defaults(params, user_training_keys)
     return params
diff --git a/tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py b/tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py
new file mode 100644
index 00000000..00488350
--- /dev/null
+++ b/tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py
@@ -0,0 +1,19 @@
+import torch
+
+
+def apply_hardware_defaults(params, explicitly_set: set) -> None:
+    """Auto-enable compile_model and native_amp when CUDA is available.
+
+    Skips fields present in explicitly_set — those are deliberate user
+    choices from the YAML config and must not be overridden.
+    hasattr guards keep this safe for params that don't carry these fields
+    yet (vision, audio — Phase 2).
+    """
+    if not torch.cuda.is_available():
+        return
+    if 'compile_model' not in explicitly_set and hasattr(params.training, 'compile_model'):
+        if getattr(params.training, 'compile_model', 0) == 0:
+            params.training.compile_model = 1
+    if 'native_amp' not in explicitly_set and hasattr(params.training, 'native_amp'):
+        if not getattr(params.training, 'native_amp', False):
+            params.training.native_amp = True
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..cc9e65bf
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_checkpoint_unwrap_compiled_model.py
@@ -0,0 +1,342 @@
+"""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.
+"""
+from argparse import Namespace
+
+import torch
+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
+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)
+
+
+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
+
+
+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)
+        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')
+    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 = _fill(_TinyModel(), 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_all_params_equal(target, 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 = _fill(_TinyModel(), 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_all_params_equal(target, 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,
+    using the same unwrap on both sides."""
+    import tempfile
+    import os
+
+    source = _fill(_TinyModel(), 2.71)
+    compiled_source = torch.compile(source, backend='aot_eager')
+    compiled_source(torch.rand(1, 4))
+
+    checkpoint = save_checkpoint(
+        # 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()
+    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_all_params_equal(fresh, 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(
+        # Namespace, not _FakeArgs: see the identical note in
+        # test_resume_from_checkpoint_symmetric_with_compiled_model.
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
+        args=Namespace(), 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)
+
+    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():
+    """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(
+        # Namespace, not _FakeArgs: see the identical note in
+        # test_resume_from_checkpoint_symmetric_with_compiled_model.
+        compiled_source, _FakeOptimizer(), _FakeScheduler(), epoch=3,
+        args=Namespace(), 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)
+
+    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
+
+    malicious_checkpoint = {
+        'model': _TinyModel().state_dict(),
+        'optimizer': {},
+        'lr_scheduler': {},
+        'epoch': 0,
+        'payload': _NotOnTheSafeGlobalsAllowlist(),
+    }
+
+    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_compile_skipped_under_quantization.py b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
new file mode 100644
index 00000000..40c94659
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_compile_skipped_under_quantization.py
@@ -0,0 +1,94 @@
+"""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.
+"""
+from unittest.mock import patch
+
+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)
+    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
+    # 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():
+    """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/tests/test_compile_warmup_fallback.py b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
new file mode 100644
index 00000000..54a26785
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_compile_warmup_fallback.py
@@ -0,0 +1,127 @@
+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():
+    """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 _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")
+
+    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)
+
+
+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/tests/test_export_model_unwrap.py b/tinyml-tinyverse/tests/test_export_model_unwrap.py
new file mode 100644
index 00000000..10792e01
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_export_model_unwrap.py
@@ -0,0 +1,101 @@
+"""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, unwrap_compiled_submodules
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+
+    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."""
+    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'))
+
+
+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/tests/test_load_weights_orig_mod.py b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
new file mode 100644
index 00000000..4276d917
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_load_weights_orig_mod.py
@@ -0,0 +1,103 @@
+"""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 _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."""
+    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_all_params_equal(target, 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_all_params_equal(target, 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_all_params_equal(target, 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_all_params_equal(target, 9.0)
diff --git a/tinyml-tinyverse/tests/test_train_best_epoch_bugs_timeseries.py b/tinyml-tinyverse/tests/test_train_best_epoch_bugs_timeseries.py
new file mode 100644
index 00000000..60ab1ed2
--- /dev/null
+++ b/tinyml-tinyverse/tests/test_train_best_epoch_bugs_timeseries.py
@@ -0,0 +1,195 @@
+"""Regression tests for two timeseries_classification/train.py and
+timeseries_forecasting/train.py post-loop bugs.
+
+1. timeseries_classification/train.py logged `best['f1']` under the
+   "AUC ROC Score" label instead of the `best['auc']` value that was
+   actually computed and stored in `best` on every improving epoch.
+
+2. When `for epoch in range(args.start_epoch, args.epochs)` runs zero
+   iterations -- e.g. `--resume` pointed at a checkpoint that already
+   satisfies `--epochs`, a real and expected use case for re-running only
+   the post-training export/compile steps -- the post-loop "Log best epoch
+   results" section read dict keys that are only ever assigned inside the
+   loop body, crashing:
+     - timeseries_classification/train.py: KeyError on best['predictions']
+       (best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) has
+       no 'predictions'/'ground_truth' keys until an improving epoch runs).
+       Now guards the whole block on `best['epoch'] is not None`.
+     - timeseries_forecasting/train.py: TypeError (None subscript) on
+       best_epoch_values['true_values'][:, :, idx] (true_values/predictions
+       are pre-populated with None, not omitted). Now guards the whole
+       block on `best_epoch_values['true_values'] is not None`.
+
+(timeseries_regression and timeseries_anomalydetection were already safe --
+their `best` dicts only ever hold scalars with safe initial values, so they
+have no test here.)
+
+These tests drive each script's real main() through a full but heavily
+mocked model/data pipeline (mocking every heavy helper imported from
+common/train_base.py and common/models.py) with args.start_epoch ==
+args.epochs, so the training loop body never executes -- reproducing
+exactly the "--resume to an already-completed checkpoint" scenario -- and
+assert main() completes without raising.
+"""
+import os
+import tempfile
+from argparse import Namespace
+from contextlib import ExitStack
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import torch
+
+from tinyml_tinyverse.references.timeseries_classification import train as tsc_train
+from tinyml_tinyverse.references.timeseries_forecasting import train as forecast_train
+
+
+class _FakeTSClassificationDataset:
+    classes = ["a", "b"]
+    X = np.zeros((4, 3, 4), dtype=np.float32)
+    inverse_label_map = {0: "a", 1: "b"}
+
+
+def _base_classification_args(**overrides):
+    args = Namespace(
+        quantization=True, dont_train_just_feat_ext='False', load_saved_model='None',
+        nas_enabled='False', generic_model=False, nn_for_feature_extraction=False,
+        output_int=True, auto_quantization=False, distributed=False, gen_golden_vectors=False,
+        model='dummy', model_config=None, model_spec=None, dual_op=False,
+        label_smoothing=0.0, apex=False, print_freq=10, quantization_method='QAT',
+        epochs=3, start_epoch=3,  # loop runs zero iterations, like a fully-resumed checkpoint
+        output_dir='/tmp/fake-output', weight_bitwidth=8, activation_bitwidth=8,
+        autoquant_tolerance_classification=0.1, opset_version=17, device='cpu',
+        file_level_classification_log='/tmp/fake-output/file_level.log', DEBUG=False,
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+def _fake_data_loaders():
+    item = (torch.tensor(0), torch.zeros(1, 3, 4), torch.tensor(0))
+    return [item], [item]
+
+
+def _set_dataset_load_state(dataset):
+    tsc_train.dataset_load_state['dataset'] = dataset
+    tsc_train.dataset_load_state['dataset_test'] = dataset
+    tsc_train.dataset_load_state['train_sampler'] = None
+    tsc_train.dataset_load_state['test_sampler'] = None
+
+
+def _patch_common_pipeline(stack):
+    """Mocks every heavy helper main() calls before/around the epoch loop,
+    common to both classification tests below."""
+    stack.enter_context(patch.object(
+        tsc_train, "setup_training_environment",
+        return_value=(tsc_train.getLogger("test"), torch.device("cpu"))))
+    stack.enter_context(patch.object(tsc_train, "prepare_transforms"))
+    stack.enter_context(patch.object(tsc_train, "create_data_loaders", return_value=_fake_data_loaders()))
+    stack.enter_context(patch.object(tsc_train.models, "get_model", return_value=torch.nn.Identity()))
+    stack.enter_context(patch.object(tsc_train, "load_pretrained_weights", side_effect=lambda model, a, l: model))
+    stack.enter_context(patch.object(tsc_train, "handle_export_only", return_value=False))
+    stack.enter_context(patch.object(tsc_train, "move_model_to_device"))
+    stack.enter_context(patch.object(tsc_train, "compile_model_if_enabled", side_effect=lambda model, a, l, **kw: model))
+    stack.enter_context(patch.object(
+        tsc_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None)))
+    stack.enter_context(patch.object(tsc_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock())))
+    stack.enter_context(patch.object(tsc_train, "resume_from_checkpoint"))
+    stack.enter_context(patch.object(tsc_train, "get_amp_context", return_value=MagicMock()))
+    stack.enter_context(patch.object(tsc_train, "get_grad_scaler", return_value=None))
+    stack.enter_context(patch.object(tsc_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model))
+    stack.enter_context(patch.object(tsc_train.utils, "export_model"))
+    stack.enter_context(patch.object(tsc_train, "log_training_time"))
+    stack.enter_context(patch.object(tsc_train, "shutdown_data_loaders"))
+
+
+def test_classification_train_main_survives_zero_iteration_resume():
+    args = _base_classification_args()
+    _set_dataset_load_state(_FakeTSClassificationDataset())
+
+    with ExitStack() as stack:
+        _patch_common_pipeline(stack)
+        tsc_train.main(0, args)
+
+
+def test_classification_train_logs_auc_not_f1_as_auc_roc_score(caplog):
+    """Runs the loop for exactly one (improving) epoch so best['auc'] and
+    best['f1'] are set to distinct values, then asserts the "AUC ROC Score"
+    log line reports the auc value, not the f1 value."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        args = _base_classification_args(
+            epochs=1, start_epoch=0, output_dir=tmp_dir,
+            file_level_classification_log=os.path.join(tmp_dir, "file_level.log"),
+        )
+        _set_dataset_load_state(_FakeTSClassificationDataset())
+
+        avg_conf_matrix = [[2, 0], [0, 2]]
+        fake_predictions = torch.zeros(4)
+        fake_ground_truth = torch.zeros(4)
+        # f1 and auc are deliberately far apart so a mislabeled log line is unmistakable.
+        evaluate_return = (99.0, 11.0, 77.0, avg_conf_matrix, fake_predictions, fake_ground_truth)
+
+        with ExitStack() as stack:
+            _patch_common_pipeline(stack)
+            stack.enter_context(patch.object(tsc_train.utils, "train_one_epoch_classification"))
+            stack.enter_context(patch.object(tsc_train.utils, "evaluate_classification", return_value=evaluate_return))
+            stack.enter_context(patch.object(tsc_train, "save_checkpoint", return_value={}))
+            stack.enter_context(patch.object(tsc_train.utils, "save_on_master"))
+            stack.enter_context(patch.object(tsc_train.utils, "print_file_level_classification_summary"))
+            with caplog.at_level("INFO"):
+                tsc_train.main(0, args)
+
+    auc_lines = [r.message for r in caplog.records if "AUC ROC Score" in r.message]
+    assert auc_lines, "expected an 'AUC ROC Score' log line"
+    assert "77.000" in auc_lines[0], f"expected the auc value (77.0) in: {auc_lines[0]}"
+    assert "11.000" not in auc_lines[0], f"AUC ROC Score line wrongly logged the f1 value: {auc_lines[0]}"
+
+
+class _FakeForecastingDataset:
+    X = np.zeros((4, 3, 4), dtype=np.float32)
+    Y = None
+    header_row = [["temp"]]
+
+
+def test_forecasting_train_main_survives_zero_iteration_resume():
+    """timeseries_forecasting/train.py crashed with TypeError (None subscript)
+    rather than KeyError, since best_epoch_values pre-populates 'true_values'
+    and 'predictions' with None instead of omitting the keys."""
+    args = Namespace(
+        quantization=True, dont_train_just_feat_ext='False', gen_golden_vectors=False,
+        model='dummy', model_config=None, model_spec=None, dual_op=False,
+        forecast_horizon=2, auto_quantization=False, distributed=False,
+        weight_bitwidth=8, activation_bitwidth=8, autoquant_tolerance_forecasting=0.1,
+        epochs=3, start_epoch=3,  # loop runs zero iterations, like a fully-resumed checkpoint
+        output_dir='/tmp/fake-output', opset_version=17, device='cpu',
+        quantization_method='QAT', output_int=None, generic_model=False,
+    )
+    dataset = _FakeForecastingDataset()
+    forecast_train.dataset_load_state['dataset'] = dataset
+    forecast_train.dataset_load_state['dataset_test'] = dataset
+    forecast_train.dataset_load_state['train_sampler'] = None
+    forecast_train.dataset_load_state['test_sampler'] = None
+
+    with ExitStack() as stack:
+        stack.enter_context(patch.object(
+            forecast_train, "setup_training_environment",
+            return_value=(forecast_train.getLogger("test"), torch.device("cpu"))))
+        stack.enter_context(patch.object(forecast_train, "prepare_transforms"))
+        stack.enter_context(patch.object(forecast_train, "generate_golden_vector_dir"))
+        stack.enter_context(patch.object(forecast_train, "create_data_loaders", return_value=_fake_data_loaders()))
+        stack.enter_context(patch.object(forecast_train.models, "get_model", return_value=torch.nn.Identity()))
+        stack.enter_context(patch.object(forecast_train, "log_model_summary"))
+        stack.enter_context(patch.object(forecast_train, "load_pretrained_weights", side_effect=lambda model, a, l: model))
+        stack.enter_context(patch.object(forecast_train, "handle_export_only", return_value=False))
+        stack.enter_context(patch.object(forecast_train, "move_model_to_device"))
+        stack.enter_context(patch.object(forecast_train, "compile_model_if_enabled", side_effect=lambda model, a, l, **kw: model))
+        stack.enter_context(patch.object(forecast_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model))
+        stack.enter_context(patch.object(forecast_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock())))
+        stack.enter_context(patch.object(
+            forecast_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None)))
+        stack.enter_context(patch.object(forecast_train, "resume_from_checkpoint"))
+        stack.enter_context(patch.object(forecast_train, "export_trained_model"))
+        stack.enter_context(patch.object(forecast_train, "log_training_time"))
+        stack.enter_context(patch.object(forecast_train, "shutdown_data_loaders"))
+        forecast_train.main(0, args)
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
index c1c04dbe..26eef23a 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py
@@ -93,13 +93,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 {}
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
index 8c2e7dee..951fb132 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
@@ -1648,6 +1648,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")
@@ -1659,6 +1677,15 @@ 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. 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:
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py
index c6a5a087..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,13 +604,41 @@ 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'])
+        # 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:
-            model_ema.load_state_dict(checkpoint['model_ema'])
+            _load_symmetric(model_ema, checkpoint['model_ema'])
     return args
 
 
@@ -658,21 +686,57 @@ 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.
+
+    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 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
     """
+    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'
@@ -684,10 +748,31 @@ 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()
+                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, 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
 
 
@@ -754,15 +839,33 @@ 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()
+        # 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
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..b55e3707 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)
@@ -361,24 +361,28 @@ def main(gpu, args):
 
         # Log best epoch results
         logger = getLogger(f"root.main.{phase}.BestEpoch")
-        logger.info("")
-        logger.info("Printing statistics of best epoch:")
-        logger.info(f"Best Epoch: {best['epoch']}")
-        logger.info(f"Acc@1 {best['accuracy']:.3f}")
-        logger.info(f"F1-Score {best['f1']:.3f}")
-        logger.info(f"AUC ROC Score {best['f1']:.3f}")
-        logger.info("")
-        logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'],
-                      columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()],
-                      index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]),
-                                                             headers="keys", tablefmt='grid')))
-
-        Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG,
-               name="root.utils.print_file_level_classification_summary",
-               append_log=True if args.quantization else False, console_log=False)
-        getLogger("root.utils.print_file_level_classification_summary").propagate = False
-        utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase)
-        logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}")
+        if best['epoch'] is not None:
+            logger.info("")
+            logger.info("Printing statistics of best epoch:")
+            logger.info(f"Best Epoch: {best['epoch']}")
+            logger.info(f"Acc@1 {best['accuracy']:.3f}")
+            logger.info(f"F1-Score {best['f1']:.3f}")
+            logger.info(f"AUC ROC Score {best['auc']:.3f}")
+            logger.info("")
+            logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'],
+                          columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()],
+                          index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]),
+                                                                 headers="keys", tablefmt='grid')))
+
+            Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG,
+                   name="root.utils.print_file_level_classification_summary",
+                   append_log=True if args.quantization else False, console_log=False)
+            getLogger("root.utils.print_file_level_classification_summary").propagate = False
+            utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase)
+            logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}")
+        else:
+            logger.warning("No epoch was run in this invocation (e.g. --resume to a checkpoint that already "
+                            "satisfied --epochs); skipping best-epoch and file-level classification summaries.")
 
         # Export model
         logger.info('Exporting model after training.')
diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py
index d7172be1..eef00c14 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
@@ -283,21 +283,25 @@ def main(gpu, args):
 
         # Log best epoch metrics
         logger = getLogger(f"root.main.{phase}.BestEpoch")
-        logger.info("Printing statistics of best epoch:")
-        logger.info(f"Best epoch:{best_epoch_values['epoch'] + 1}")
-        logger.info(f"Overall SMAPE across all variables: {best_epoch_values['overall_smape']:.2f}%")
-        logger.info("Per-Variable Metrics:")
-
-        for idx, item in enumerate(dataset.header_row):
-            for target_variable_name in item:
-                logger.info(f"  Variable {target_variable_name}:")
-                logger.info(f"      SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(best_epoch_values['true_values'][:, :, idx], best_epoch_values['predictions'][:, :, idx]):.2f}%")
-                logger.info(f"      R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(best_epoch_values['predictions'][:, :, idx], best_epoch_values['true_values'][:, :, idx]):.4f}")
-
-                for step in range(args.forecast_horizon):
-                    logger.info(f"      Timestep {step + 1}:")
-                    logger.info(f"          SMAPE: {utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]):.2f}%")
-                    logger.info(f"          R²: {utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]):.4f}")
+        if best_epoch_values['true_values'] is not None:
+            logger.info("Printing statistics of best epoch:")
+            logger.info(f"Best epoch:{best_epoch_values['epoch'] + 1}")
+            logger.info(f"Overall SMAPE across all variables: {best_epoch_values['overall_smape']:.2f}%")
+            logger.info("Per-Variable Metrics:")
+
+            for idx, item in enumerate(dataset.header_row):
+                for target_variable_name in item:
+                    logger.info(f"  Variable {target_variable_name}:")
+                    logger.info(f"      SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(best_epoch_values['true_values'][:, :, idx], best_epoch_values['predictions'][:, :, idx]):.2f}%")
+                    logger.info(f"      R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(best_epoch_values['predictions'][:, :, idx], best_epoch_values['true_values'][:, :, idx]):.4f}")
+
+                    for step in range(args.forecast_horizon):
+                        logger.info(f"      Timestep {step + 1}:")
+                        logger.info(f"          SMAPE: {utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]):.2f}%")
+                        logger.info(f"          R²: {utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]):.4f}")
+        else:
+            logger.warning("No epoch was run in this invocation (e.g. --resume to a checkpoint that already "
+                            "satisfied --epochs); skipping best-epoch metrics summary.")
 
         # Save final predictions and create visualizations for best epoch
         if args.output_dir and best_epoch_values['true_values'] is not None:
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