add div prints

ACEsuit · Oct 1, 2024 · 24bf8a6 · 24bf8a6
1 parent 814e734
commit 24bf8a6
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/mace/modules/blocks.py b/mace/modules/blocks.py
@@ -667,9 +667,12 @@ def forward(
         )  # [n_nodes, irreps]
         print("message before linear", message[1, :10])
         print("message before linear", message[3, :10])
+        print("self.avg_num_neighbors", self.avg_num_neighbors)
+        print("message before div", self.linear(message)[1, :10])
+        print("message before div", self.linear(message)[3, :10])
         message = self.linear(message) / self.avg_num_neighbors
-        print("message", message[1, :10])
-        print("message", message[3, :10])
+        print("message with div", message[1, :10])
+        print("message with div", message[3, :10])
         return (
             self.reshape(message),
             sc,