diff --git a/src/components/entity/ha-entity-attribute-picker.ts b/src/components/entity/ha-entity-attribute-picker.ts
index 8e8ffb2adb..a4d46c3605 100644
--- a/src/components/entity/ha-entity-attribute-picker.ts
+++ b/src/components/entity/ha-entity-attribute-picker.ts
@@ -17,6 +17,7 @@ import {
 } from "lit-element";
 import memoizeOne from "memoize-one";
 import { fireEvent } from "../../common/dom/fire_event";
+import { computeDomain } from "../../common/entity/compute_domain";
 import { PolymerChangedEvent } from "../../polymer-types";
 import { HomeAssistant } from "../../types";
 import { formatAttributeName } from "../../util/hass-attributes-util";
@@ -179,9 +180,7 @@ class HaEntityAttributePicker extends LitElement {
     }
 
     return Object.keys(stateObj.attributes).filter((attr) =>
-      SELECTABLE_ATTRIBUTES[entity.substring(0, entity.indexOf("."))].includes(
-        attr
-      )
+      SELECTABLE_ATTRIBUTES[computeDomain(entity)].includes(attr)
     );
   });