diff --git a/src/hotspot/share/gc/shenandoah/shenandoahUnload.cpp b/src/hotspot/share/gc/shenandoah/shenandoahUnload.cpp
index 8ff9ff2b876..bb13e9b8e22 100644
--- a/src/hotspot/share/gc/shenandoah/shenandoahUnload.cpp
+++ b/src/hotspot/share/gc/shenandoah/shenandoahUnload.cpp
@@ -50,8 +50,7 @@ class ShenandoahIsUnloadingOopClosure : public OopClosure {
 
 public:
   ShenandoahIsUnloadingOopClosure() :
-    // TODO: In non-generational mode, this should still be complete_marking_context()
-    _marking_context(ShenandoahHeap::heap()->marking_context()),
+    _marking_context(ShenandoahHeap::heap()->complete_marking_context()),
     _is_unloading(false) {
   }