len(x) -> x.size(0)

mir-group · Oct 25, 2024 · 77a0e9f · 77a0e9f
1 parent 5f1c9e8
commit 77a0e9f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nequip/nn/_interaction_block.py b/nequip/nn/_interaction_block.py
@@ -174,7 +174,7 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         if avg_num_neigh is not None:
             edge_features = edge_features.div(avg_num_neigh**0.5)
         # now scatter down
-        x = scatter(edge_features, edge_dst, dim=0, dim_size=len(x))
+        x = scatter(edge_features, edge_dst, dim=0, dim_size=x.size(0))
 
         x = self.linear_2(x)