diff --git a/features/notify_helpers.py b/features/notify_helpers.py
index 5f77307..f22ed0e 100644
--- a/features/notify_helpers.py
+++ b/features/notify_helpers.py
@@ -123,8 +123,7 @@
   return notified
 
 
-# TODO: change to FlaskInternalTask when convert to flask
-class NotifyTaskBase(jsonfeed.InternalTask):
+class NotifyTaskBase(jsonfeed.FlaskInternalTask):
   """Abstract base class for notification task handler."""
 
   _EMAIL_TEMPLATE = None  # Subclasses must override this.