add pos to gearnet required_batch_attributes #73

a-r-j · Feb 7, 2024 · 81be059 · 81be059
1 parent 5572403
commit 81be059
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/proteinworkshop/models/graph_encoders/gear_net.py b/proteinworkshop/models/graph_encoders/gear_net.py
@@ -125,6 +125,7 @@ def required_batch_attributes(self) -> Set[str]:
         """
         return {
             "x",
+            "pos",
             "edge_index",
             "edge_type",
             "edge_attr",