fixed bug with precomputed feature input

Author: aimat

examples/nn_pes_butene.py

@@ -98,14 +98,14 @@
 }
 
 nn.create({
-    # 'e': hyper_energy,
-    'eg': hyper_grads,
+    'e': hyper_energy,
+    # 'eg': hyper_grads,
     # 'nac2': hyper_nac2
 })
 
 y = {
-    # 'e': eng,
-    'eg': [eng, grads],
+    'e': eng,
+    # 'eg': [eng, grads],
     # 'nac2' : nac,
 }
 

pyNNsMD/models/mlp_e.py

@@ -152,7 +152,8 @@ def call(self, data, training=False, **kwargs):
             temp_g = tape2.batch_jacobian(temp_e, x)
             y_pred = [temp_e, temp_g]
         elif self.precomputed_features:
-            feat_flat_std = self.std_layer(x)
+            x1 = x[0]
+            feat_flat_std = self.std_layer(x1)
             temp_hidden = self.mlp_layer(feat_flat_std, training=training)
             temp_e = self.energy_layer(temp_hidden)
             y_pred = temp_e

pyNNsMD/nn_pes_src/training/training_mlp_e.py

@@ -166,14 +166,14 @@ def train_model_energy(i=0, outdir=None, mode='training'):
     x_rescale, y1 = scaler.transform(x, y)
 
     # Model + Model precompute layer +feat
-    feat_x, _ = out_model.precompute_feature_in_chunks(x_rescale, batch_size=batch_size)
+    feat_x, feat_grad = out_model.precompute_feature_in_chunks(x_rescale, batch_size=batch_size)
     feat_x_mean, feat_x_std = out_model.set_const_normalization_from_features(feat_x,normalization_mode=normalize_feat)
 
 
     # Train Test split
-    xtrain = feat_x[i_train]
+    xtrain = [feat_x[i_train], feat_grad[i_train]]
     ytrain = y1[i_train]
-    xval = feat_x[i_val]
+    xval = [feat_x[i_val], feat_grad[i_val]]
     yval = y1[i_val]
 
     # Compile model