Optimize convolution and groupnorm (#98)

* Add a channels-last-native Triton GroupNorm

ScaFFold's GroupNorm was the only operator forcing a layout change: ATen's
kernel launches one workgroup per (batch, group) row -- 8 of them at the
benchmark's defaults -- so on a 228-CU MI300A it ran at a small fraction of
achievable bandwidth, and it was the sole reason channels-last broke.

This is an NDHWC-native kernel: Welford statistics, fp32/bf16/fp16 with
torch's autocast contract, an int64 tile-base path for tensors past 2^31
elements, and a fused ReLU that is bit-exact against the unfused form.
Registered as a custom op with fake kernels and register_autograd, so it
traces under torch.compile(fullgraph=True) and composes with DistConv's
DCTensor.

Determinism is a contract here, not an accident: no float atomics, and the
grid, split, tile and reduction order are pure functions of the shape.
Verified run-to-run and across interpreters.

Nothing imports this yet -- the wiring is the next commit.

* Route the UNet GroupNorm through the Triton kernel

FastGroupNorm becomes a three-rung ladder -- the native Triton kernel, a
torch.compile'd functional, and stock eager -- with the routing decision
made per call and the rejections tested in order: an explicit opt-out, a
rung that has already failed in this process, an active torch.func
transform, a non-CUDA tensor, a tensor subclass other than DistConv's
DCTensor, a GPU the launch tables were not tuned on, and anything the
kernel's own is_supported rejects.

The ladder primitives live in _rungs.py, which the convolution ladder will
share. The hardware guard is a preference rather than a correctness
condition -- the kernel is correct anywhere Triton lowers it, and what is
unknown elsewhere is only its speed -- so an explicit opt-in overrides it
and nothing else on that list can be overridden at all.

DoubleConv now asks the GroupNorm for its ReLU: the Triton kernel folds the
activation into its forward store, removing a streaming pass worth 38% of
the forward at the shapes that dominate the step. The nn.ReLU slots stay
occupied by nn.Identity so nn.Sequential does not renumber its children and
existing checkpoints keep loading.

* Add a Triton 3-D convolution package

A self-contained NDHWC implicit-GEMM convolution for MI300A / gfx942, built
to replace MIOpen in this benchmark and to be upstreamable to DistConv on
its own terms. It imports nothing from ScaFFold.

Four kernels serve seven operator-directions. Backward-data is not a kernel
at all -- it is the forward contraction on a permuted weight -- and the
transposed operator's backward directions are its own forward kernel and
backward-weight with the operands swapped. That reuse is the main
structural result.

The corpus is recorded from real ScaFFold calls, and it records the *shape
form*: the logical convolution, the halo'd unpadded one DistConv hands a
backend, and the padded one this package's adapter actually issues. Those
are different problems, and conflating them has been the most expensive
mistake in this work, so ConvProblem names which is which and the benchmark
driver takes --form.

Backward-weight is deterministic by default: split-K with a reduction tree
whose split count, tile and order are pure functions of the shape.
Correctness is checked against fp64 references under a three-tier tolerance
policy, with a bitwise corpus for the exactly-representable cases.

Also here: the benchmark harness, which times kernels through CUDA-graph
replay with 95% intervals and an online iteration count, and which can run
without the MIOpen control -- 98% of a two-arm capture's wall clock was
MIOpen's find, not measurement.

* Route every convolution through the Triton rung

FastConv3d and FastConvTranspose3d mirror FastGroupNorm: drop-in nn.Conv3d
and nn.ConvTranspose3d, same parameters under the same names, no buffers, a
rung ladder sharing _rungs.py, and MIOpen underneath everything the kernel
declines.

The transposed operator gets a factory of its own rather than a flag,
because it is a different operator -- the weight's channel axes are the
other way round and a different set of kernels sits behind it.

The adapter performs the halo exchange itself, above autograd, rather than
leaving it to the one DistConv does below. That has a consequence for the
shape the kernel sees, and it is the thing to know when reading any number
from this work: only the split axis is halo'd, so padding=1 survives on the
other two and every k=3 convolution in the network is padded at every
configuration -- unsharded there is nothing to halo at all. The benchmark
corpus calls that the adapter form, and it is what production issues.

* Cite findings rather than the scratch tree that produced them

Every comment and docstring in the committed tree stood on its own path into
the untracked work/ scratch directory: 70 references across 20 files, all of
them dangling for anyone who clones this branch.  The measurement each one
supported is kept and stated as a result -- the number, the direction of the
effect, and the "this was tried and it loses" warnings that stop a closed
question being re-litigated -- while the path, the capture filename, the
section number of an unshipped document and the blow-by-blow methodology go.

Three pointers were dead even with work/ present and are simply gone: the
transposed benchmark driver the conv_bench docstring narrated a refactor away
from, a tuned table's named source capture, and the review commit SHA in the
GroupNorm wiring tests.

The corpus JSON keeps model-analysis/unet_shapes.py as its provenance; nothing
parses either file's "source" field.

* Add helper script to run tests.

* Satisfy the style checks CI actually runs

The style workflow runs `ruff format --diff .` and `ruff check .`, and both
were failing: 22 files would be reformatted and there were 19 lint errors,
all of the latter in the Triton package.  Most of this is whitespace.

Three fixes are not mechanical:

`gemm_probe` built two closures over `a` and `b` in functions whose
`finally` deletes both names, so each was correct only because the harness
happens to call it before the cleanup runs.  They now bind the tensors as
default arguments, which captures at definition time and does not depend
on call order.  This is also what ruff was reporting as F821.

`baseline._callable` assigned two lambdas to names, now plain functions.

Six imports were unused and are gone.

`ScaFFold/viz/standard_viz.py` is not part of this branch's work -- it
fails the format check on round2-fixes too, and is reformatted here only
so the check can pass.

Both suites re-run afterwards, since several tests read their subject's
source: triton_conv3d 1425 passed / 16 skipped, ScaFFold 742 passed /
8 skipped / 1 xfailed.  Reformatting changes Triton's JIT cache key, so
the first run after this recompiles every kernel and takes ~5x longer.

* Say which kernel each accelerated module ended up using

A run gives no sign of whether the Triton rungs actually served it.  Print
one line per ladder on rank 0 at startup, naming Triton against everything
PyTorch does -- compiled and eager together, since from outside the ladder
those are the same answer.

Placement is the whole of the problem.  ``_triton_ok`` is a latch set when a
rung first answers a call, so reporting at construction time would say
"Native" about modules that have not run yet.  Reporting at the end of
warmup is right when there is a warmup, but ``warmup_batches <= 0`` returns
before running anything, which is what the benchmark drivers configure -- so
a single call site would drop the line on exactly the runs most likely to
want it.  Both sites call it behind a one-shot flag.

The reporter is duck-typed on the latch rather than on isinstance: _rungs is
imported by the modules it would otherwise have to import back, and a ladder
added later is then reported without touching this code.  Labels come from a
_rung_label class attribute and fall back to the class name, so a ladder that
does not declare one still appears.

Deliberately one rank's answer and not a collective: ranks latch
independently and a rung can still fall back later, so this is informational
rather than a contract, and gathering it would put a barrier on a path with
no other reason for one.

* Tighten comments and drop measured figures from them

Reviewers asked that the comments on this branch be more concise and that
specific performance numbers be removed from comments describing changes,
since those numbers go stale.  Comments and docstrings only: every file's
AST is unchanged apart from docstrings, ruff format and check are clean,
and the test runner's shell script differs only in comment lines.

Gone: timings, speedups, roofline percentages, memory figures, cell
tallies, dated history of what the code used to do, and pointers to the
untracked scratch tree.  Kept: contracts, traps, and the hardware and
algorithmic constants the code depends on (buffer-op and int32 limits,
tile geometry, tolerances, pinned test shapes).  Where a rationale would
be empty without its figure, the comment names the bench tool that
regenerates it.

Across the 30 files: 11.9k comment and docstring lines become 10.6k, and
703 lines carrying a measurement, date, or scratch-tree pointer become 65,
each of which is a constraint rather than a measurement.

* Install the triton_conv3d package alongside ScaFFold

ScaFFold.unet.conv3d imports the sibling package, but package discovery
only listed ScaFFold*, so a wheel shipped an adapter with nothing behind
it.  Discovery is namespace-aware, so the test tree (and the recorded
corpora it holds) is excluded by name; the suite runs from a checkout.
Verified by building a wheel, installing it into an empty target and
importing both packages from there.

* Keep the recorded convolution corpora with the tests

Nothing on the runtime path reads scaffold_corpus.json or
scaffold_census.json; only bench/ and the test suite do, and both run
from a checkout.  Move them under triton_conv3d/tests/data, point the
loaders in shapes.py there, and say what each file is a record of: the
corpus joins model-analysis shape dumps with MIOpen profiles of the same
configurations, the census is an instrumented training run.

* Say why triton_group_norm's activation defaults to None

ScaFFold always fuses the ReLU; the default exists because the function
is a stand-alone replacement for F.group_norm.

* Never take the Triton conv rung without a handler for its failures

_triton_kernel_failures() resolved to () when triton.errors could not be
imported, and `except ()` catches nothing, so on a Triton without that
module every handler in the ladder became a bare try and the first
CompilationError or OutOfResources took the rank down -- the inverse of
the ladder's contract, silently.  Resolve the 2.x roots
(CompilationError, OutOfResources) when triton.errors is absent, and have
_routing_declines refuse the rung, with one warning, when nothing at all
resolved.

* Agree the halo exchange across the mesh instead of deciding it per rank

Every input to the Triton rung's routing verdict is a fact about one
rank -- the thin-shard check in _halo_plan, the device, the layout, the
shape gates, the latch -- while the halo exchange is a collective.  A rank
declining alone went through DistConv's exchange while its neighbour
posted receives against this module's, and the mesh hung.  Unreachable
with the shipped shapes, where D divides evenly, but a hang.

The global extent is not on a DCTensor, so the check cannot be derived
locally; instead, for any call some rank may exchange on, whether to
exchange through this module is the MIN of every rank's verdict, agreed
once per module instance and cached.  The latch, and any later change in
a rank's own verdict, only choose the kernel run on the exchanged tensor,
which communicates nothing.  A shard that later becomes too thin for a
plan on an agreed module is an error rather than a lone trip to MIOpen.

Adds the first test of the real exchange: two ranks under torchrun, the
sharded output checked against F.conv3d on the whole volume, then a local
decline, a latch and a thin shard, each ending with both ranks on the
same side of the exchange.

* Commit the generator behind scaffold_corpus.json

make_corpus.py traces the three profiled configurations with
model-analysis/unet_shapes.py, the pure-stdlib shape calculator the corpus
already cited as its source, and joins profile_points.json, the MIOpen
profile record of the same runs reduced to its convolution kernels.  Both
were in the untracked scratch tree; the corpus was not regenerable from
the repository.  `--check` confirms the committed file is reproduced
byte for byte.

* Commit the instrument and generator behind scaffold_census.json

conv_census.py wraps FastConv3d, the two autograd nodes and the six kernel
entry points around a real `scaffold benchmark` run and writes one capture
per rank; make_census.py folds the four captures under census/ into the
committed file.  Both were in the untracked scratch tree, where the
instrument ran through a scratch training driver.  Verified by running
the new instrument at configuration A: its capture records the same
modules, autograd nodes and kernel calls as the committed one on every
field except the resolver's answers on padded backward-weight problems,
which the tuning tables have moved since and the census does not read.
`--check` confirms the committed census is reproduced from the captures.

The census's stale `tuned_declined` field is dropped from the instrument:
the decline it described no longer exists.  A README in tests/data says
what each file is and how to regenerate it.

Author: ndryden

ScaFFold/unet/_rungs.py

@@ -20,6 +20,7 @@
 
 from ScaFFold.utils.perf_measure import annotate
 
+from .conv3d import FastConv3d, FastConvTranspose3d
 from .group_norm import FastGroupNorm
 
 _doubleconv_annotate = annotate(fmt="DoubleConv.{}")
@@ -28,30 +29,126 @@
 _outconv_annotate = annotate(fmt="OutConv.{}")
 
 
-def _group_norm(num_groups, num_channels):
+def _group_norm(num_groups, num_channels, activation=None):
     if num_channels % num_groups != 0:
         raise ValueError(
             f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}"
         )
-    # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the
-    # same parameters under the same names, so checkpoints are unaffected.
-    return FastGroupNorm(num_groups, num_channels)
+    # FastGroupNorm is nn.GroupNorm plus a Triton/compiled GPU kernel; it holds
+    # the same parameters under the same names, and `activation` is a plain
+    # attribute rather than a submodule, so checkpoints are unaffected.
+    return FastGroupNorm(num_groups, num_channels, activation=activation)
+
+
+def _conv3d(in_channels, out_channels, **kwargs):
+    """The model's non-transposed convolutions, in one place.
+
+    ``FastConv3d`` is ``nn.Conv3d`` plus a Triton GPU kernel; it holds the same
+    parameters under the same names and adds no buffers, so checkpoints are
+    unaffected in either direction.  It falls back to MIOpen for anything the
+    kernel does not serve, and it serves the sharded configurations too: it
+    performs the halo exchange itself, above autograd, rather than leaving it to
+    the one DistConv does below.  Only the split axis is halo'd, so ``padding=1``
+    survives on the other two and these convolutions are padded at every
+    configuration; see :mod:`ScaFFold.unet.conv3d`.
+
+    The ``nn.ConvTranspose3d`` in ``Up`` do not come through here: they are a
+    different operator, with the weight's channel axes the other way round and a
+    different set of kernels behind them, so they have a factory of their own
+    (:func:`_conv_transpose3d`).
+    """
+    return FastConv3d(in_channels, out_channels, **kwargs)
+
+
+def _conv_transpose3d(in_channels, out_channels, **kwargs):
+    """The model's transposed convolutions -- the decoder's upsamplers.
+
+    ``FastConvTranspose3d`` is ``nn.ConvTranspose3d`` plus a Triton GPU kernel;
+    it holds the same parameters (``weight`` *and* ``bias``, which these sites
+    have and the ordinary convolutions mostly do not) under the same names and
+    adds no buffers, so checkpoints are unaffected in either direction.  The
+    kernel serves only the ``kernel == stride``, no-padding upsample built
+    below; anything else falls back to MIOpen.
+    """
+    return FastConvTranspose3d(in_channels, out_channels, **kwargs)
+
+
+def _consumer_dtype(*tensors):
+    """The dtype the convolution consuming a concatenation will actually see.
+
+    Inside an enabled autocast region the answer is autocast's dtype, because
+    ``aten::convolution`` carries the ``lower_precision_fp`` cast policy and
+    casts whatever it is handed.  Producing that dtype from the concatenation is
+    *bitwise identical* to producing ATen's promoted dtype and letting the
+    convolution narrow it -- the promoted tensor holds exact widenings of both
+    sources, so narrowing before or after the copy rounds the same values once
+    -- while moving half the bytes.
+
+    Outside autocast the answer is ``torch.cat``'s ordinary promotion, so eval,
+    ``inference_mode`` and pure-fp32 runs are unchanged.
+    """
+    dtype = tensors[0].dtype
+    for tensor in tensors[1:]:
+        dtype = torch.promote_types(dtype, tensor.dtype)
+    device_type = tensors[0].device.type
+    try:
+        if not torch.is_autocast_enabled(device_type):
+            return dtype
+        autocast_dtype = torch.get_autocast_dtype(device_type)
+    except (RuntimeError, TypeError):  # a device type autocast does not know
+        return dtype
+    # Only ever narrow: if autocast's dtype is the wider of the two, keep the
+    # promotion ATen would have done.
+    if torch.promote_types(autocast_dtype, dtype) is autocast_dtype:
+        return dtype
+    return autocast_dtype
+
+
+def _skip_concat(skip, upsampled):
+    """``torch.cat([skip, upsampled], dim=1)`` at the consumer's dtype.
+
+    Under ``torch.autocast`` the two halves do not share a dtype: the skip comes
+    from a GroupNorm, an fp32-policy op, while the upsampled half comes from a
+    ``ConvTranspose3d`` and is bf16.  ``torch.cat`` carries the ``promote``
+    policy, so it widens the bf16 half to fp32, concatenates at fp32, and the
+    following convolution narrows the whole double-width result straight back
+    down -- three full-resolution passes to deliver one.  Casting the inputs
+    first collapses that to one, and the convolution reads the same bits either
+    way (see :func:`_consumer_dtype`).
+    """
+    dtype = _consumer_dtype(skip, upsampled)
+    return torch.cat([skip.to(dtype), upsampled.to(dtype)], dim=1)
 
 
 class DoubleConv(nn.Module):
-    """(convolution => GroupNorm => ReLU) * 2"""
+    """(convolution => GroupNorm => ReLU) * 2
+
+    The ReLU lives *inside* the GroupNorm (``activation="relu"``), because the
+    Triton GroupNorm kernel folds it into its forward store for free and thereby
+    removes an entire streaming pass.  ``FastGroupNorm`` applies the ReLU on
+    every path, including eager, so the network's function is unchanged; only
+    the number of memory passes differs.
+
+    The ``nn.ReLU`` slots are held open by ``nn.Identity`` rather than removed:
+    ``nn.Sequential`` names its children by position, so deleting them would
+    renumber the two convolutions and the second GroupNorm and invalidate every
+    existing checkpoint.  Neither ``nn.ReLU`` nor ``nn.Identity`` has parameters
+    or buffers, so with the placeholders in place the state dict is byte
+    identical to the pre-fusion model's (pinned by
+    ``tests/test_groupnorm.py::test_state_dict_matches_plain_groupnorm_model``).
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=None):
         super().__init__()
         if not mid_channels:
             mid_channels = out_channels
         self.double_conv = nn.Sequential(
-            nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, mid_channels),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, out_channels),
-            nn.ReLU(inplace=True),
+            _conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            _group_norm(group_norm_groups, mid_channels, activation="relu"),
+            nn.Identity(),
+            _conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            _group_norm(group_norm_groups, out_channels, activation="relu"),
+            nn.Identity(),
         )
 
     @_doubleconv_annotate
@@ -75,7 +172,16 @@ def forward(self, x):
 
 
 class Up(nn.Module):
-    """Upscaling then double conv"""
+    """Upscaling then double conv
+
+    The skip concatenation goes through :func:`_skip_concat` rather than
+    ``torch.cat`` directly, so that it emits the dtype the following
+    convolution will use instead of ``torch.cat``'s promoted one.  The tensor
+    that convolution reads is bitwise unchanged either way; it is written and
+    read back at half the width, saving both time and peak memory.  This rests
+    on ``self.conv`` beginning with a convolution, which the constructor below
+    guarantees on either branch.
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True):
         super().__init__()
@@ -90,7 +196,7 @@ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True)
                 in_channels // 2,
             )
         else:
-            self.up = nn.ConvTranspose3d(
+            self.up = _conv_transpose3d(
                 in_channels, in_channels // 2, kernel_size=2, stride=2
             )
             self.conv = DoubleConv(in_channels, out_channels, group_norm_groups)
@@ -118,14 +224,16 @@ def forward(self, x1, x2):
         # if you have padding issues, see
         # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
         # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
-        x = torch.cat([x2, x1], dim=1)
+        # torch.cat([x2, x1], dim=1) at the dtype the convolution below
+        # actually wants; see the class docstring.
+        x = _skip_concat(x2, x1)
         return self.conv(x)
 
 
 class OutConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(OutConv, self).__init__()
-        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+        self.conv = _conv3d(in_channels, out_channels, kernel_size=1)
 
     @_outconv_annotate
     def forward(self, x):

ScaFFold/unet/conv3d.py

@@ -29,6 +29,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from ScaFFold.unet._rungs import format_kernel_selection, kernel_selection
 from ScaFFold.utils.checkpointing import CheckpointManager
 from ScaFFold.utils.data_loading import FractalDataset, SpatialShardSpec
 from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_TORCH_DTYPE
@@ -90,6 +91,9 @@ class BaseTrainer:
 
     def __init__(self, model, config, device, log):
         self.model = model
+        # One-shot guard for the startup kernel-selection line; see
+        # _log_kernel_selection for why it has two call sites.
+        self._kernel_selection_logged = False
         self.config = config
         self.device = device
         self.log = log
@@ -786,6 +790,30 @@ def _warmup_ragged_batches(self, batch):
                 log_prefix=f"warmup ragged ({ragged}): ",
             )
 
+    def _log_kernel_selection(self):
+        """Log which kernel each accelerated module is using, once, on rank 0.
+
+        Called from two places -- the end of :meth:`warmup` and after the first
+        training batch -- because the answer only exists once a forward has run:
+        ``_triton_ok`` is a latch set when a rung first answers a call, so
+        reporting at construction time would say "Native" about modules that
+        have simply not run yet.  ``warmup_batches <= 0`` makes :meth:`warmup`
+        return without running anything, so neither call site alone covers every
+        run; the flag below makes the pair idempotent.
+
+        Rank 0 only, and one rank's answer: each rank latches independently, so
+        under DDP this is representative rather than global.  Deliberately not a
+        collective -- gathering it would put a barrier on a path that has no
+        other reason for one.
+        """
+        if self._kernel_selection_logged or self.world_rank != 0:
+            return
+        self._kernel_selection_logged = True
+        model = getattr(self.model, "module", self.model)
+        self.log.info("Kernel selection (rank 0):")
+        for line in format_kernel_selection(kernel_selection(model)):
+            self.log.info(line)
+
     def warmup(self):
         """Run warmup iterations before the main training loop."""
         warmup_batches = self.config.warmup_batches
@@ -847,6 +875,7 @@ def warmup(self):
 
         torch.distributed.barrier()
         self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s")
+        self._log_kernel_selection()
 
     def train(self, profiler=None):
         """
@@ -953,6 +982,8 @@ def train(self, profiler=None):
                         # not skew the epoch mean.
                         train_dice_total += batch_dice_score * batch_size
                         end_code_region("run_training_batch")
+                        if first_batch:
+                            self._log_kernel_selection()
 
                         # Update the loss
                         begin_code_region("update_loss")

ScaFFold/unet/group_norm.py

@@ -49,7 +49,9 @@ def main(config: RunConfig):
                 if "val_loss_avg" in row:
                     val_loss.append(float(row["val_loss_avg"]))
 
-        plot_title = f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}"
+        plot_title = (
+            f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}"
+        )
         line_thickness = 2
         fontsize = 20
         tick_fontsize = 14
@@ -64,7 +66,9 @@ def main(config: RunConfig):
         plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
         plt.yscale("log")
         plt.title(plot_title, fontsize=12)
-        plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+        plt.legend(
+            loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+        )
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight")
         plt.close(figures[-1])
@@ -76,7 +80,9 @@ def main(config: RunConfig):
         plt.ylabel("Val dice score", fontsize=fontsize)
         plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
         plt.title(plot_title, fontsize=12)
-        plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+        plt.legend(
+            loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+        )
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight")
         plt.close(figures[-1])
@@ -89,7 +95,9 @@ def main(config: RunConfig):
             plt.ylabel("Val loss", fontsize=fontsize)
             plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
             plt.title(plot_title, fontsize=12)
-            plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+            plt.legend(
+                loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+            )
             plt.grid(True, axis="y")
             plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight")
             plt.close(figures[-1])