--- a
+++ b/tests/training/test_train.py
@@ -0,0 +1,209 @@
+# ruff:noqa:E402
+
+import pytest
+
+try:
+    import torch.nn
+except ImportError:
+    torch = None
+
+if torch is None:
+    pytest.skip("torch not installed", allow_module_level=True)
+pytest.importorskip("rich")
+
+import shutil
+from typing import (
+    Optional,
+    Sequence,
+    Union,
+)
+
+import pytest
+import spacy.tokenizer
+import torch.nn
+from confit import Config
+from confit.utils.random import set_seed
+from spacy.tokens import Doc, Span
+
+from edsnlp.core.registries import registry
+from edsnlp.data.converters import AttributesMappingArg, get_current_tokenizer
+from edsnlp.metrics.dep_parsing import DependencyParsingMetric
+from edsnlp.training.optimizer import LinearSchedule, ScheduledOptimizer
+from edsnlp.training.trainer import GenericScorer, train
+from edsnlp.utils.span_getters import SpanSetterArg, set_spans
+
+
+@registry.factory.register("myproject.custom_dict2doc", spacy_compatible=False)
+class CustomSampleGenerator:
+    def __init__(
+        self,
+        *,
+        tokenizer: Optional[spacy.tokenizer.Tokenizer] = None,
+        name: str = "myproject.custom_dict2doc",
+        span_setter: SpanSetterArg = {"ents": True, "*": True},
+        bool_attributes: Union[str, Sequence[str]] = [],
+        span_attributes: Optional[AttributesMappingArg] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.name = name
+        self.span_setter = span_setter
+        self.bool_attributes = bool_attributes
+        self.span_attributes = span_attributes
+
+    def __call__(self, obj):
+        tok = get_current_tokenizer() if self.tokenizer is None else self.tokenizer
+        doc = tok(obj["note_text"] or "")
+        doc._.note_id = obj.get("note_id", obj.get("__FILENAME__"))
+        doc._.note_datetime = obj.get("note_datetime")
+
+        spans = []
+
+        if self.span_attributes is not None:
+            for dst in self.span_attributes.values():
+                if not Span.has_extension(dst):
+                    Span.set_extension(dst, default=None)
+
+        for ent in obj.get("entities") or ():
+            ent = dict(ent)
+            span = doc.char_span(
+                ent.pop("start"),
+                ent.pop("end"),
+                label=ent.pop("label"),
+                alignment_mode="expand",
+            )
+            for label, value in ent.items():
+                new_name = (
+                    self.span_attributes.get(label, None)
+                    if self.span_attributes is not None
+                    else label
+                )
+                if self.span_attributes is None and not Span.has_extension(new_name):
+                    Span.set_extension(new_name, default=None)
+
+                if new_name:
+                    span._.set(new_name, value)
+            spans.append(span)
+
+        set_spans(doc, spans, span_setter=self.span_setter)
+        for attr in self.bool_attributes:
+            for span in spans:
+                if span._.get(attr) is None:
+                    span._.set(attr, False)
+        return doc
+
+
+def test_ner_qualif_train_diff_bert(run_in_test_dir, tmp_path):
+    set_seed(42)
+    config = Config.from_disk("ner_qlf_diff_bert_config.yml")
+    shutil.rmtree(tmp_path, ignore_errors=True)
+    kwargs = Config.resolve(config["train"], registry=registry, root=config)
+    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)
+    scorer = GenericScorer(**kwargs["scorer"])
+    val_data = kwargs["val_data"]
+    last_scores = scorer(nlp, val_data)
+
+    # Check empty doc
+    nlp("")
+
+    assert last_scores["ner"]["micro"]["f"] > 0.4
+    assert last_scores["qual"]["micro"]["f"] > 0.4
+
+
+def test_ner_qualif_train_same_bert(run_in_test_dir, tmp_path):
+    set_seed(42)
+    config = Config.from_disk("ner_qlf_same_bert_config.yml")
+    shutil.rmtree(tmp_path, ignore_errors=True)
+    kwargs = Config.resolve(config["train"], registry=registry, root=config)
+    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)
+    scorer = GenericScorer(**kwargs["scorer"])
+    val_data = kwargs["val_data"]
+    last_scores = scorer(nlp, val_data)
+
+    # Check empty doc
+    nlp("")
+
+    assert last_scores["ner"]["micro"]["f"] > 0.4
+    assert last_scores["qual"]["micro"]["f"] > 0.4
+
+
+def test_qualif_train(run_in_test_dir, tmp_path):
+    set_seed(42)
+    config = Config.from_disk("qlf_config.yml")
+    shutil.rmtree(tmp_path, ignore_errors=True)
+    kwargs = Config.resolve(config["train"], registry=registry, root=config)
+    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)
+    scorer = GenericScorer(**kwargs["scorer"])
+    val_data = kwargs["val_data"]
+    last_scores = scorer(nlp, val_data)
+
+    # Check empty doc
+    nlp("")
+
+    assert last_scores["qual"]["micro"]["f"] >= 0.4
+
+
+def test_dep_parser_train(run_in_test_dir, tmp_path):
+    set_seed(42)
+    config = Config.from_disk("dep_parser_config.yml")
+    shutil.rmtree(tmp_path, ignore_errors=True)
+    kwargs = Config.resolve(config["train"], registry=registry, root=config)
+    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)
+    scorer = GenericScorer(**kwargs["scorer"])
+    val_data = list(kwargs["val_data"])
+    last_scores = scorer(nlp, val_data)
+
+    scorer_bis = GenericScorer(parser=DependencyParsingMetric(filter_expr="False"))
+    # Just to test what happens if the scores indicate 2 roots
+    val_data_bis = [Doc.from_docs([val_data[0], val_data[0]])]
+    nlp.pipes.parser.decoding_mode = "mst"
+    last_scores_bis = scorer_bis(nlp, val_data_bis)
+    assert last_scores_bis["parser"]["uas"] == 0.0
+
+    # Check empty doc
+    nlp("")
+
+    assert last_scores["dep"]["las"] >= 0.4
+
+
+def test_optimizer():
+    net = torch.nn.Linear(10, 10)
+    optim = ScheduledOptimizer(
+        torch.optim.AdamW,
+        module=net,
+        total_steps=10,
+        groups={
+            ".*": {
+                "lr": 9e-4,
+                "schedules": LinearSchedule(
+                    warmup_rate=0.1,
+                    start_value=0,
+                ),
+            }
+        },
+    )
+    for param in net.parameters():
+        assert "exp_avg" not in optim.optim.state[param]
+    optim.initialize()
+    for param in net.parameters():
+        assert "exp_avg" in optim.optim.state[param]
+    lr_values = [optim.optim.param_groups[0]["lr"]]
+    for i in range(10):
+        optim.step()
+        lr_values.append(optim.optim.param_groups[0]["lr"])
+
+    # close enough
+    assert lr_values == pytest.approx(
+        [
+            0.0,
+            0.0009,
+            0.0008,
+            0.0007,
+            0.0006,
+            0.0005,
+            0.0004,
+            0.0003,
+            0.0002,
+            0.0001,
+            0.0,
+        ]
+    )