[lsbharadwaj]

Summary

Fix GPU training (#231). The device mismatch ("input layer has device = cpu and the core has device = cpu") came from FLAMO comparing the string "cpu" with the torch.device that .to() stores; FLAMO 0.2.20 (gdalsanto/flamo#38) canonicalizes module devices, so this PR requires flamo>=0.2.20. On the pyFDN side:

• model_response excites the model on the model's own device.
• wrap_fdn_shell / dss_to_flamo / trainable_from_build build the FFT layers on the requested device; wrap_fdn_shell defaults to the core's own device.
• train_fdn defaults to the model's device instead of "cpu", so a CUDA model is not moved back to the CPU for training.
• Examples pass device="cpu" explicitly where they build FLAMO modules by hand.

Fixes #231

Testing

• pytest passes locally (523 passed, 1 skipped on FLAMO 0.2.20; CPU only)
• New tests added for new functionality: Shell from a device="cpu" core plus .to("cpu") layers (fails on 0.2.19); build/render/train and wrap_fdn_shell stay on the model's device, on CPU and on CUDA/MPS when available (skipped without a GPU)
• Coverage not decreased (pytest --cov=pyFDN)
• Training example run on a molab GPU (GPU training still faulty #231)

Checklist

• Pre-commit hooks pass (pre-commit run --all-files)
• Docstrings updated for any changed public API
• HISTORY.rst updated if this is a user-facing change

[lsbharadwaj]

@SebastianJiroSchlecht Kindly have a look at the fix

[SebastianJiroSchlecht]

Review of 38cd49a (pass torch.device instead of "cpu")

Diagnosis. FLAMO's Shell checks at construction that core, input and output layer share one .device, comparing with == — and torch.device("cpu") == "cpu" is False. Hence the confusing "input layer has device = cpu and the core has device = cpu". The two types arise because pyFDN's builders store whatever the caller passed (the string "cpu"), while since FLAMO 0.2.19 any .to(...) call resets .device to the torch.device object (e.g. dsp.FFT(nfft).to("cpu") in the allpass example).

This commit avoids the mismatch in one example and in coupled_rooms, but every other caller passing device="cpu" (~15 examples, and users) can still hit it as soon as a .to() is mixed in.

FLAMO's set_device. I tested it (FLAMO 0.2.19): model.set_device(dev) / .to(dev) moves the whole graph and re-syncs every module's .device. But it only helps after a Shell exists — a mismatch already fails inside the Shell constructor.

Suggested fix

1. Build on CPU and call model.set_device(device) once at the end of wrap_fdn_shell / dss_to_flamo / trainable_from_build, instead of threading device through every module builder.
2. Normalising strings belongs in FLAMO, not pyFDN — I'm opening a FLAMO PR for that, so "cpu" and torch.device("cpu") compare equal everywhere. After it's released, the example edits in this commit aren't needed.
3. Tests: one CPU test that assembles a Shell from parts built with device="cpu" plus .to("cpu") (reproduces the bug), and one CUDA test skipped when no GPU is available.

Also: #231 is "GPU training doesn't run on molab", and CI is CPU-only — please run the training example on a molab GPU before we merge.

[SebastianJiroSchlecht]

Follow-up: switching to torch.device doesn't fix this on a GPU either. torch.device("cuda") == torch.device("cuda:0") is False, and .to("cuda") stores cuda:0, so any layer moved with .to() still mismatches a core built with torch.device("cuda"). I reproduced this with mps, which behaves the same way. That may well be the molab failure in #231.

Fix upstream: gdalsanto/flamo#38 stores every module's device the way a tensor reports it. With that branch, pyFDN's full suite passes with the string device="cpu" examples unchanged, so once it's released this PR can drop the example edits and bump flamo>=.

[SebastianJiroSchlecht]

Pushed a commit: FLAMO 0.2.20 includes the device fix (gdalsanto/flamo#38), so this PR now requires flamo>=0.2.20 and reverts 38cd49a. The allpass example is back on device="cpu", and test_marimo_examples_run covers it as the regression test. The full suite passes locally on FLAMO 0.2.20 (412 passed, 1 skipped). The example also picked up a ruff-format reflow from the pre-commit hook.

Still open before merging: a run of the training example on a molab GPU (#231). The optional set_device refactor can go in a separate PR.

[SebastianJiroSchlecht]

Pushed afafdbc (merge of main) and 9966d1f:

• train_fdn defaulted to "cpu" while the builders default to CUDA when one is available, so training with default arguments moved a GPU model back to the CPU. It now defaults to model.device.
• wrap_fdn_shell(core, nfft=…) with no device built its output layer on CUDA even when the core was on the CPU, and Shell's device check then failed. It now follows the core's device.
• Restored model.set_nfft(train_nfft) / set_nfft(render_nfft) in example_train_fdn_to_rir: 567bb30 had removed them, so the example trained on the render grid.
• Fixed the misindented device="cpu" lines in example_multislope_rir_to_fdn.
• Added the tests from the review:
  • a CPU regression test that builds a Shell from a device="cpu" core plus .to("cpu") layers. I checked it fails on FLAMO 0.2.19 and passes on 0.2.20.
  • build, render and train with default arguments stay on the model's device, and wrap_fdn_shell keeps the core's device. These run on the CPU and, when available, on CUDA/MPS; without a GPU the CUDA cases are skipped.
• Added a HISTORY.rst entry and updated the description.

The full suite passes locally on FLAMO 0.2.20, which is the current PyPI release: 523 passed, 1 skipped.

Still open: a run of the training example on a molab GPU (#231). CI and my environment are CPU-only, so the CUDA test above has not run yet. It will run on a GPU with pytest tests/test_train.py -k device.

---

Generated by Claude Code

[lsbharadwaj]

@SebastianJiroSchlecht
I am not able to recreate the bug. I am doing the following steps and I dont see any issues

1. Use marimo to run examples/example_train_fdn_to_rir.py using marimo edit examples/example_train_fdn_to_rir.py
2. The program by default runs with GPU
3. Change to CPU using drop down. (Works)
4. Change to GPU again using drop down. (Works)

I repeated this 2 times

Please tell me if I am missing something