From 33b33be2cca1a355b5e3ee60c84baec95e17d268 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 7 Apr 2024 06:42:22 +0000 Subject: [PATCH] fix:UTs --- deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 ++-- deepmd/pt/model/model/make_model.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 5c1db5a550..6cede6a265 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -199,9 +199,9 @@ def forward_atomic( mapping, fparam, aparam, - )["energy"] + ) # apply bias to each individual model - ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype])) + ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype][:,:nloc])["energy"]) weights = self._compute_weight(extended_coord, extended_atype, nlists_) fit_ret = { diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index f1a0334335..ab3f3d0a8f 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -172,6 +172,10 @@ def forward_common( model_predict = self.output_type_cast(model_predict, input_prec) return model_predict + @property + def out_bias(self)-> torch.Tensor: + return self.atomic_model.out_bias + def change_out_bias( self, merged,