fix error in error fix

proteinworkshop/tasks/ppi_site_prediction.py

@@ -30,15 +30,14 @@ def __init__(self, radius: float = 3.5, ca_only: bool = True) -> None:
         self.fill_value = 1e-5
         self.ca_only = ca_only
         charstr: str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
-        self.chain_map: Dict[str, int] = {
-            charstr[i]: i for i in range(len(charstr))
-        }
+        self.chain_map: Dict[str, int] = {charstr[i]: i for i in range(len(charstr))}
 
     def __call__(self, data: Union[Protein, Data]):
         # Map the chain labels to integers
         target_chains = []
 
         chain_strs = [res.split(":")[0] for res in data.residue_id]
+        chain_strs = np.unique(chain_strs)
 
         for chain in data.graph_y:
             target_chains.append(chain_strs.index(chain))
@@ -59,9 +58,7 @@ def __call__(self, data: Union[Protein, Data]):
         # Unwrap the coordinates
         other_chains = other_chains.reshape(-1, 3)
         # Remove any rows with 1e-5
-        other_chains = other_chains[
-            ~torch.all(other_chains == self.fill_value, dim=1)
-        ]
+        other_chains = other_chains[~torch.all(other_chains == self.fill_value, dim=1)]
 
         # Create a KDTree
         # If Ca only, we only see if the interacting chains are within the