implemented feedback

datamol-io · Oct 25, 2023 · ee56497 · ee56497
1 parent 8ceb167
commit ee56497
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/graphium/nn/pyg_layers/gps_pyg.py b/graphium/nn/pyg_layers/gps_pyg.py
@@ -253,9 +253,9 @@ def forward(self, batch: Batch) -> Batch:
         if self.node_residual:
             if self.layer_depth < 1:
                 h_local = self.residual_add(h_local, feat_in)
-                h_local *= self.scale_activations(h_local, self.output_scale)
+                h_local = self.scale_activations(h_local, self.output_scale)
             else:
-                h_local *= self.scale_activations(h_local, self.output_scale)
+                h_local = self.scale_activations(h_local, self.output_scale)
                 h_local = self.residual_add(h_local, feat_in)
 
         if self.norm_layer_local is not None: