Merge pull request #80 from linusyh/fix-sequence-edges

Fix `sequence_edges` when b is a Data object
a-r-j · Mar 2, 2024 · 3bea793 · 3bea793
2 parents e6aee4f + 76ab97b
commit 3bea793
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,8 @@
 ### Command
 * Adds `--force-cuda-version` to `workshop install` [#78](https://github.com/a-r-j/ProteinWorkshop/pull/78)
 
+### Features
+* Fix `sequence_edges` behaviour when argument `b` is a `Data` object [#80](https://github.com/a-r-j/ProteinWorkshop/pull/80)
 
 ### 0.2.5 (28/12/2023)
 

diff --git a/proteinworkshop/features/edges.py b/proteinworkshop/features/edges.py
@@ -104,7 +104,7 @@ def sequence_edges(
         idx_b = torch.arange(1, b.ptr[-1], device=b.ptr.device)
     elif isinstance(b, Data):
         idx_a = torch.arange(0, b.coords.shape[0] - 1, device=b.coords.device)
-        idx_a = torch.arange(1, b.coords.shape[0] - 1, device=b.coords.device)
+        idx_b = torch.arange(1, b.coords.shape[0], device=b.coords.device)
     # Concatenate indices to create edge list
     if direction == "forward":
         e_index = torch.stack([idx_a, idx_b], dim=0)