fix(pt): pass device and dtype in SeZMDeNSFittingNet.deserialize

[Anai-Guo]

Summary

SeZMDeNSFittingNet.deserialize cannot run — it raises

TypeError: safe_numpy_to_tensor() missing 2 required keyword-only arguments: 'device' and 'dtype'

safe_numpy_to_tensor is declared in deepmd/pt/model/descriptor/sezm_nn/utils.py as

def safe_numpy_to_tensor(
    data: Any, *, device: torch.device, dtype: torch.dtype
) -> torch.Tensor:

Both are keyword-only and required, but dens.py:749 called it with the array alone:

state = {key: safe_numpy_to_tensor(value) for key, value in variables.items()}

Why this is the right fix

This is the only one of the 17 safe_numpy_to_tensor call sites in the tree that omits them. The other sezm_nn deserialize implementations — ffn.py, norm.py, embedding.py, block.py, attn_res.py, activation.py, radial.py, so2.py, so3.py — all share one identical block, so I simply matched it:

obj = cls(**config)
template = obj.state_dict()
state = {
    key: safe_numpy_to_tensor(
        value, device=template[key].device, dtype=template[key].dtype
    )
    for key, value in variables.items()
}
obj.load_state_dict(state)
return obj

Taking device/dtype from the freshly-built module's own state_dict() is what makes the round trip self-consistent: serialize builds @variables from state = self.state_dict(), so every key in variables is present in template and template[key] cannot KeyError.

I deliberately did not copy the siblings' precision = config.pop("precision"); config["dtype"] = PRECISION_DICT[precision] lines — SeZMDeNSFittingNet.__init__ takes precision directly and derives self.dtype itself, so cls(**config) is already correct here.

Verification

I don't have a runtime for this path, so I verified by replaying the real signature (AST-extracted from utils.py, not hand-copied) against the real call site:

DEF: def safe_numpy_to_tensor(data, *, device, dtype)
CURRENT dens.py call -> TypeError: missing a required keyword-only argument: 'device'
FIXED   call         -> BINDS OK

ruff check and ruff format --check (v0.16.0, the pinned pre-commit version) pass on the changed file.

---

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved model deserialization to preserve each parameter’s device and data type, ensuring more reliable loading across different hardware and precision settings.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: 8732a77f-8434-4613-8af8-4bfe42a9ca29

📥 Commits

Reviewing files that changed from the base of the PR and between be1a636 and b37ddb8.

📒 Files selected for processing (1)

• deepmd/pt/model/descriptor/sezm_nn/dens.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

The deserialization path now converts serialized tensors to each reconstructed parameter’s device and dtype before loading the state dictionary.

Device-aware state loading

Layer / File(s)	Summary
Device-aware state loading deepmd/pt/model/descriptor/sezm_nn/dens.py	deserialize uses template state entries to preserve the target device and dtype when converting serialized tensors.
Estimated code review effort: 2 (Simple)	~10 minutes

Merge Risk: ⚪ Minimal · up to b37dd

This change makes deserialization load serialized tensors using reconstructed parameter devices and dtypes. No current merge-blocking risk is identified.

Suggested reviewers: outisli

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the fix to SeZMDeNSFittingNet.deserialize and the required device and dtype arguments. It matches the main change.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.96%. Comparing base (be1a636) to head (b37ddb8).

Files with missing lines	Patch %	Lines
deepmd/pt/model/descriptor/sezm_nn/dens.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6006      +/-   ##
==========================================
- Coverage   77.20%   76.96%   -0.25%     
==========================================
  Files        1152     1152              
  Lines      138792   138793       +1     
  Branches     5062     5056       -6     
==========================================
- Hits       107156   106818     -338     
- Misses      29754    30093     +339     
  Partials     1882     1882              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.