diff --git a/module/ui/draggable_list.py b/module/ui/draggable_list.py
index 806ff35b0..7f8f46c92 100644
--- a/module/ui/draggable_list.py
+++ b/module/ui/draggable_list.py
@@ -42,11 +42,9 @@ class DraggableList:
         self.name = name
         self.keyword_class = keyword_class
         self.ocr_class = ocr_class
-        if not isinstance(keyword_class, list):
-            keyword_class = [keyword_class]
-        self.known_rows = [
-            kw for kc in keyword_class for kw in kc.instances.values()
-        ]
+        if isinstance(keyword_class, list):
+            keyword_class = keyword_class[0]
+        self.known_rows = list(keyword_class.instances.values())
         self.search_button = search_button
         self.check_row_order = check_row_order
         self.active_color = active_color
diff --git a/tasks/assignment/ui.py b/tasks/assignment/ui.py
index 8aef740a0..d184c2d5b 100644
--- a/tasks/assignment/ui.py
+++ b/tasks/assignment/ui.py
@@ -121,6 +121,10 @@ ASSIGNMENT_ENTRY_LIST = DraggableList(
     check_row_order=False,
     active_color=(40, 40, 40)
 )
+ASSIGNMENT_ENTRY_LIST.known_rows = [
+    kw for kc in ASSIGNMENT_ENTRY_LIST.keyword_class
+    for kw in kc.instances.values()
+]
 
 
 class AssignmentUI(UI):