fix(pt): implement get_task_dim method in GeneralFitting base class for proper inheritance

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

Author: Copilot

checkpoint

@@ -126,9 +126,3 @@ def forward_lower(
         else:
             model_predict = model_ret
         return model_predict
-
-    @torch.jit.export
-    def get_task_dim(self) -> int:
-        """Get the output dimension of the dipole model."""
-        # For dipole models, the output dimension is always 3 (x, y, z components)
-        return 3

deepmd/pt/model/model/dipole_model.py

@@ -657,3 +657,8 @@ def _forward_common(
         outs = torch.where(mask[:, :, None], outs, 0.0)
         results.update({self.var_name: outs})
         return results
+
+    @torch.jit.export
+    def get_task_dim(self) -> int:
+        """Get the output dimension of the fitting net."""
+        return self._net_out_dim()

deepmd/pt/model/task/fitting.py

@@ -53,8 +53,7 @@ class TestInferDeepTensorPt : public ::testing::Test {
   deepmd::DeepTensor dt;
 
   void SetUp() override {
-    std::string file_name =
-        "../../tests/infer/deepdipole_pt_with_get_task_dim.pth";
+    std::string file_name = "../../tests/infer/deepdipole_pt.pth";
     dt.init(file_name);
   };