From 5d2a3a9314b495d2f3088cf044ffbe0a466ff5f8 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 14 Feb 2024 22:17:17 +0800 Subject: [PATCH 1/2] fix bug of output dimension --- deepmd/pt/model/task/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 1b3e2c3d65..2cd58ccc42 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -119,7 +119,7 @@ def __init__( self.aparam_avg, self.aparam_inv_std = None, None in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - out_dim = 1 + out_dim = self.dim_out self.old_impl = kwargs.get("old_impl", False) if self.old_impl: @@ -358,7 +358,7 @@ def forward( dim=-1, ) - outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion + outs = torch.zeros(nf, nloc, self.dim_out) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None From cc7d8b0de56ffa59c62bba7866bada8015715245 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 14 Feb 2024 22:23:49 +0800 Subject: [PATCH 2/2] rm out_dim --- deepmd/pt/model/task/ener.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 2cd58ccc42..81f2cc8cf0 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -119,7 +119,6 @@ def __init__( self.aparam_avg, self.aparam_inv_std = None, None in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - out_dim = self.dim_out self.old_impl = kwargs.get("old_impl", False) if self.old_impl: @@ -144,7 +143,7 @@ def __init__( networks=[ FittingNet( in_dim, - out_dim, + self.dim_out, self.neuron, self.activation_function, self.resnet_dt,