diff --git a/src/driver.ts b/src/driver.ts
index 564f55b..3a197f6 100644
--- a/src/driver.ts
+++ b/src/driver.ts
@@ -8,7 +8,7 @@ import { getState, resetState, setState } from "./state";
 import "./driver.css";
 
 export type DriveStep = {
-  element?: string | Element;
+  element?: string | Element | (() => Element);
   onHighlightStarted?: DriverHook;
   onHighlighted?: DriverHook;
   onDeselected?: DriverHook;
diff --git a/src/highlight.ts b/src/highlight.ts
index 961a26a..214023a 100644
--- a/src/highlight.ts
+++ b/src/highlight.ts
@@ -29,7 +29,8 @@ function mountDummyElement(): Element {
 
 export function highlight(step: DriveStep) {
   const { element } = step;
-  let elemObj = typeof element === "string" ? document.querySelector(element) : element;
+  let elemObj =
+    typeof element === "function" ? element() : typeof element === "string" ? document.querySelector(element) : element;
 
   // If the element is not found, we mount a 1px div
   // at the center of the screen to highlight and show