Remove undirected edge transform

JustinDoIt · Sep 12, 2023 · 59716b9 · 59716b9
1 parent 4946a58
commit 59716b9
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 5 deletions.
diff --git a/proteinworkshop/features/edges.py b/proteinworkshop/features/edges.py
@@ -9,7 +9,6 @@
 from graphein.protein.tensor.data import Protein, ProteinBatch
 from omegaconf import ListConfig
 from torch_geometric.data import Batch, Data
-from torch_geometric.utils import to_undirected
 
 
 @beartype
@@ -59,7 +58,7 @@ def compute_edges(
     for edge_type in edge_types:
         if edge_type.startswith("knn") or edge_type.startswith("eps"):
             edges.append(
-                to_undirected(edge_fn(x.pos, edge_type))
+                edge_fn(x.pos, edge_type)
             )
         elif edge_type == "seq_forward":
             edges.append(

diff --git a/proteinworkshop/models/graph_encoders/tfn.py b/proteinworkshop/models/graph_encoders/tfn.py
@@ -171,9 +171,6 @@ def forward(self, batch: Union[Batch, ProteinBatch]) -> EncoderOutput:
             the dimension of the embeddings.
         :rtype: EncoderOutput
         """
-        # from torch_geometric.utils import to_undirected
-        # edge_index = to_undirected(batch.edge_index)
-
         # Node embedding
         h = self.emb_in(batch.x)  # (n,) -> (n, d)