From 46822efde879b46d99f2d36e5a3ff9964bbc0da4 Mon Sep 17 00:00:00 2001
From: Liza Kozlova <el.mi.kozlova@gmail.com>
Date: Thu, 28 Dec 2023 16:02:17 +0000
Subject: [PATCH] docs: update docstring

---
 proteinflow/data/torch.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/proteinflow/data/torch.py b/proteinflow/data/torch.py
index 28fbf59..8b003c2 100644
--- a/proteinflow/data/torch.py
+++ b/proteinflow/data/torch.py
@@ -908,8 +908,9 @@ def set_cdr(self, cdr):
 
         Parameters
         ----------
-        cdr : {"H1", "H2", "H3", "L1", "L2", "L3"}
-            The CDR to be iterated over. Set to `None` to go back to iterating over all chains.
+        cdr : list | str | None
+            The CDR to be iterated over (choose from H1, H2, H3, L1, L2, L3).
+            Set to `None` to go back to iterating over all chains.
 
         """
         if not self.sabdab: