fix: cdr bug

adaptyvbio · Dec 28, 2023 · 29f6eb5 · 29f6eb5
1 parent 46822ef
commit 29f6eb5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/proteinflow/data/torch.py b/proteinflow/data/torch.py
@@ -1064,7 +1064,7 @@ def __getitem__(self, idx):
  id = self.data[idx] # data is already filtered by length
  chain_id = random.choice(list(self.files[id].keys()))
  if self.cdr is not None:
- while self.cdr != chain_id.split("__")[1]:
+ while chain_id.split("__")[1] not in self.cdr:
  chain_id = random.choice(list(self.files[id].keys()))
  else:
  cluster = self.data[idx]