diff --git a/src/askui/locators/locators.py b/src/askui/locators/locators.py
index 7e3f3870..652c08b0 100644
--- a/src/askui/locators/locators.py
+++ b/src/askui/locators/locators.py
@@ -13,8 +13,8 @@
 """The type of match to use.
 
 - `"similar"` uses a similarity threshold to determine if the text is a match.
-- `"exact"` requires the text to be exactly the same (this is not the same as `"similar"` 
-  with a `similarity_threshold` of `100` as a `similarity_threshold` of `100` can still 
+- `"exact"` requires the text to be exactly the same (this is not the same as `"similar"`
+  with a `similarity_threshold` of `100` as a `similarity_threshold` of `100` can still
   allow for small differences in very long texts).
 - `"contains"` requires the text to contain (exactly) the specified text.
 - `"regex"` uses a regular expression to match the text.
@@ -70,7 +70,7 @@ class Element(Locator):
     """Locator for finding ui elements by their class.
 
     Args:
-        class_name (Literal["text", "textfield"] | None, optional): The class of the ui element, e.g., `'text'` or `'textfield'`. Defaults to `None`.
+        class_name (Literal["switch","text", "textfield"] | None, optional): The class of the ui element, e.g., `'text'` or `'textfield'`. Defaults to `None`.
 
     Examples:
         ```python
@@ -87,10 +87,7 @@ class Element(Locator):
     @validate_call
     def __init__(
         self,
-        class_name: Annotated[
-            Literal["text", "textfield"] | None,
-            Field(),
-        ] = None,
+        class_name: Literal["switch", "text", "textfield"] | None = None,
     ) -> None:
         super().__init__()
         self._class_name = class_name
diff --git a/tests/e2e/agent/test_locate.py b/tests/e2e/agent/test_locate.py
index 317451f4..04bf14b0 100644
--- a/tests/e2e/agent/test_locate.py
+++ b/tests/e2e/agent/test_locate.py
@@ -7,12 +7,7 @@
 
 from askui.agent import VisionAgent
 from askui.exceptions import ElementNotFoundError
-from askui.locators import (
-    AiElement,
-    Element,
-    Prompt,
-    Text,
-)
+from askui.locators import AiElement, Element, Prompt, Text
 from askui.locators.locators import Image
 from askui.models import ModelName
 
@@ -39,6 +34,20 @@ def test_locate_with_string_locator(
         assert 450 <= x <= 570
         assert 190 <= y <= 260
 
+    def test_locate_with_switch_class_locator(
+        self,
+        vision_agent: VisionAgent,
+        path_fixtures: pathlib.Path,
+        model: str,
+    ) -> None:
+        """Test locating elements using a class locator."""
+        locator = Element("switch")
+        x, y = vision_agent.locate(
+            locator, path_fixtures / "images" / "switch.png", model=model
+        )
+        assert 340 <= x <= 360
+        assert 420 <= y <= 460
+
     def test_locate_with_textfield_class_locator(
         self,
         vision_agent: VisionAgent,
diff --git a/tests/fixtures/images/switch.png b/tests/fixtures/images/switch.png
new file mode 100644
index 00000000..aa92424c
Binary files /dev/null and b/tests/fixtures/images/switch.png differ