docs(pt_expt): improve torch_module docstring and update tests

- Enhanced torch_module decorator docstring to document auto-generation behavior
- Updated tests to use module(...) instead of module.forward(...) to test the full
  invocation path through torch.nn.Module.__call__
- This ensures tests cover hooks, tracing, and export behavior properly

Addresses review feedback from PR #5246

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/pt_expt/common.py

@@ -300,6 +300,20 @@ def torch_module(
 ) -> type[torch.nn.Module]:
     """Convert a NativeOP to a torch.nn.Module.
 
+    This decorator wraps a NativeOP class to make it a PyTorch module, handling
+    initialization, attribute setting, and method delegation automatically.
+
+    **Auto-generated methods:**
+
+    - If the wrapped class has a ``call()`` method but does not explicitly define
+      ``forward()``, a ``forward()`` method will be auto-generated that delegates
+      to ``call()``.
+    - If the wrapped class has a ``call_lower()`` method but does not explicitly
+      define ``forward_lower()``, a ``forward_lower()`` method will be auto-generated
+      that delegates to ``call_lower()``.
+    - Explicit ``forward()`` or ``forward_lower()`` definitions in the wrapped class
+      are always respected and will not be overridden.
+
     Parameters
     ----------
     module : type[NativeOP]
@@ -308,13 +322,13 @@ def torch_module(
     Returns
     -------
     type[torch.nn.Module]
-        The torch.nn.Module.
+        The torch.nn.Module with auto-generated delegation methods if applicable.
 
     Examples
     --------
     >>> @torch_module
     ... class MyModule(NativeOP):
-    ...     pass
+    ...     pass  # forward() auto-generated from call() if it exists
     """
 
     @wraps(module, updated=())

source/tests/pt_expt/utils/test_common.py

@@ -40,7 +40,7 @@ class MockModule(MockNativeOP):
 
     module = MockModule()
     input_tensor = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cpu"))
-    output = module.forward(input_tensor)
+    output = module(input_tensor)
     expected = input_tensor * 2
     assert torch.allclose(output, expected)
 
@@ -82,7 +82,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
     module = MockModule()
     input_tensor = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cpu"))
-    output = module.forward(input_tensor)
+    output = module(input_tensor)
     expected = input_tensor * 3  # Should use the explicit forward, not call()
     assert torch.allclose(output, expected)