diff --git a/lms/djangoapps/edxnotes/helpers.py b/lms/djangoapps/edxnotes/helpers.py
index fb8de803af..c8a28b26dc 100644
--- a/lms/djangoapps/edxnotes/helpers.py
+++ b/lms/djangoapps/edxnotes/helpers.py
@@ -175,7 +175,7 @@ def preprocess_collection(user, course, collection):
 
             model.update(update)
             usage_id = model["usage_id"]
-            if usage_id in cache:
+            if usage_id in list(cache.keys()):
                 model.update(cache[usage_id])
                 filtered_collection.append(model)
                 continue
@@ -204,13 +204,13 @@ def preprocess_collection(user, course, collection):
                 if not section:
                     log.debug(u"Section not found: %s", usage_key)
                     continue
-                if section in cache:
-                    usage_context = cache[section]
+                if section.location in list(cache.keys()):
+                    usage_context = cache[section.location]
                     usage_context.update({
                         "unit": get_module_context(course, unit),
                     })
                     model.update(usage_context)
-                    cache[usage_id] = cache[unit] = usage_context
+                    cache[usage_id] = cache[unit.location] = usage_context
                     filtered_collection.append(model)
                     continue
 
@@ -218,14 +218,14 @@ def preprocess_collection(user, course, collection):
                 if not chapter:
                     log.debug(u"Chapter not found: %s", usage_key)
                     continue
-                if chapter in cache:
-                    usage_context = cache[chapter]
+                if chapter.location in list(cache.keys()):
+                    usage_context = cache[chapter.location]
                     usage_context.update({
                         "unit": get_module_context(course, unit),
                         "section": get_module_context(course, section),
                     })
                     model.update(usage_context)
-                    cache[usage_id] = cache[unit] = cache[section] = usage_context
+                    cache[usage_id] = cache[unit.location] = cache[section.location] = usage_context
                     filtered_collection.append(model)
                     continue
 
@@ -236,9 +236,9 @@ def preprocess_collection(user, course, collection):
             }
             model.update(usage_context)
             if include_path_info:
-                cache[section] = cache[chapter] = usage_context
+                cache[section.location] = cache[chapter.location] = usage_context
 
-            cache[usage_id] = cache[unit] = usage_context
+            cache[usage_id] = cache[unit.location] = usage_context
             filtered_collection.append(model)
 
     return filtered_collection