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

[vboussot]

Summary

Fixes a GPU crash for ensemble models whose custom Network.load() appends modules at load time
(e.g. TotalSegmentator, whose load() adds a Head.Conv sized from the checkpoint's class count).

The bug

Network.to (network.py) places modules by walking the tree and moving each one whose placement marker
is still gpu == "cpu" — but it only descends into a child when that child itself is unplaced. So a
module added inside an already-placed parent (Head) after the initial placement is never reached and
stays on CPU. Since ModelComposite._ensure_model_loaded re-runs the custom load() per model swap, the
freshly-added head is recreated on CPU every time, and the forward pass crashes.

The reported error depends on the path: under autocast: true (fp16) torch reports the dtype mismatch
first —

RuntimeError: Input type (c10::Half) and bias type (float) should be the same

— and in fp32 the device mismatch —

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

Both are the same root cause (a CPU-resident module receiving a GPU tensor); there is no separate
autocast bug.

The fix

After each per-model load() in ModelComposite._ensure_model_loaded, co-locate 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. The base nn.Module.parameters/buffers
are used to read each module's own tensors, since ModuleArgsDict overrides those without a recurse kwarg.

Test plan

• New tests/unit/test_model_load_device.py: a model whose load() appends a head is placed on the GPU,
  then asserted to have the head co-located (GPU-gated) + a CPU no-op guard.
• test_predictor_memory.py and test_perf_hot_paths.py (both exercise ModelComposite) still pass.
• End-to-end: TotalSegmentator-KonfAI:total-3mm now runs to completion on GPU under autocast: true
  and writes a valid whole-body segmentation (59 labels present, range 0–117). Before this fix it crashed
  on the first forward.