diff --git a/includes/flag.features.inc b/includes/flag.features.inc
index 23ad242..cf21c23 100644
--- a/includes/flag.features.inc
+++ b/includes/flag.features.inc
@@ -19,8 +19,24 @@ function flag_features_export($data, &$export, $module_name = '') {
   foreach ($data as $key => $flag) {
     // Get the module that provided the flag definition.
     if ($flag = flag_load($flag, TRUE)) {
-      $module = $modules[$flag->entity_type];
-      $export['dependencies'][$module] = $module;
+      if (array_key_exists($flag->entity_type, $modules)) {
+        $module = $modules[$flag->entity_type];
+      }
+      else {
+        // This is a flag on a custom entity. Let's figure out which module
+        // defines it.
+        if ($entity_info = entity_get_info($flag->entity_type)) {
+          if (isset($entity_info['module'])) {
+            $module = $entity_info['module'];
+          }
+        }
+      }
+
+      // Just in case.
+      if (isset($module)) {
+        $export['dependencies'][$module] = $module;
+      }
+
       $export['features']['flag'][$flag->name] = $flag->name;
     }
   }