fix(predict): co-locate load-added modules onto the model's device

konfai/predictor.py

@@ -678,6 +678,29 @@ def _predict_log(
                     )
 
 
+def _colocate_loaded_modules(model: torch.nn.Module) -> None:
+    """Move any still-CPU leaf module onto the model's device.
+
+    A custom :meth:`Network.load` may append modules after the model was already placed on its
+    device — e.g. a head sized from the checkpoint's class count — and those default to CPU, which
+    then raises a device mismatch on the forward pass. This re-homes any fully-CPU leaf onto the
+    device the rest of the model already lives on. Modules already on a device (including
+    model-parallel splits across several GPUs) are left untouched.
+    """
+    target = next((p.device for p in model.parameters() if p.device.type != "cpu"), None)
+    if target is None:
+        return
+    for sub in model.modules():
+        # ModuleArgsDict overrides parameters()/buffers() without a ``recurse`` kwarg, so use the
+        # base nn.Module methods to read each module's own (non-recursive) tensors.
+        own = [
+            *torch.nn.Module.parameters(sub, recurse=False),
+            *torch.nn.Module.buffers(sub, recurse=False),
+        ]
+        if own and all(t.device.type == "cpu" for t in own):
+            sub.to(target)
+
+
 class ModelComposite(Network):
     """
     A composite model that replicates a given base network multiple times and combines their outputs.
@@ -744,6 +767,9 @@ def _ensure_model_loaded(self, index: int) -> Network:
             # ensemble suffix added after the previous load.
             model.set_name(self._base_model_name)
             model.load(state, init=False)
+            # A custom load() may append checkpoint-sized modules (e.g. the head) on CPU; co-locate
+            # them with the already device-placed model so the forward pass doesn't hit a mismatch.
+            _colocate_loaded_modules(model)
             model.set_name(f"{self._base_model_name}_{index}")
             self._loaded_state_index = index
         return model

tests/unit/test_model_load_device.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from konfai.network.network import Network
+from konfai.predictor import Mean, ModelComposite, _colocate_loaded_modules
+
+
+class _LateHeadNetwork(Network):
+    """Model whose load() appends a checkpoint-sized module, mimicking TotalSegmentator's Head.Conv."""
+
+    def __init__(self) -> None:
+        super().__init__(in_channels=1)
+        self.add_module("Stem", torch.nn.Conv3d(1, 2, kernel_size=1))
+
+    def load(self, state_dict, init: bool = True, ema: bool = False):  # type: ignore[override]
+        # A head sized from the checkpoint, created at load time -> defaults to CPU.
+        self.add_module("Head", torch.nn.Conv3d(2, int(state_dict["nb_class"]), kernel_size=1))
+
+    def forward(self, batch_sample, output_layers=[]):  # type: ignore[override]
+        return []
+
+
+def test_colocate_is_a_safe_noop_when_model_is_all_cpu() -> None:
+    # With no device-placed parameter there is nothing to co-locate; the helper must be a no-op.
+    model = torch.nn.Sequential(torch.nn.Conv3d(1, 2, 1), torch.nn.Conv3d(2, 3, 1))
+    _colocate_loaded_modules(model)
+    assert all(not p.is_cuda for p in model.parameters())
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="device co-location only manifests on GPU")
+def test_ensemble_load_colocates_late_added_head_on_gpu() -> None:
+    # Reproduces the TotalSegmentator crash: the model is placed on the GPU, then a per-model load()
+    # appends a Head on CPU. Before the fix the forward hit "Input cuda, weight CPU".
+    composite = ModelComposite(_LateHeadNetwork(), Mean())
+    Network.to(composite, 0)  # place on cuda:0, exactly as the predictor does before inference
+
+    composite.load([{"nb_class": 5}])  # single source -> triggers _ensure_model_loaded(0)
+    model = composite["Model_0"]
+
+    head = dict(model.named_modules())["Head"]
+    assert list(head.parameters()), "test setup: Head should have parameters"
+    assert all(p.is_cuda for p in head.parameters()), "load-added Head must be co-located onto the GPU"
+    # the whole model must live on a single device
+    assert len({p.device for p in model.parameters()}) == 1