diff --git a/typedspark/_core/dataset.py b/typedspark/_core/dataset.py
index 5fcf3d6..de4b14d 100644
--- a/typedspark/_core/dataset.py
+++ b/typedspark/_core/dataset.py
@@ -186,7 +186,7 @@ def __setattr__(self, name: str, value: Any) -> None:
         if name == "__orig_class__":
             orig_class_args = get_args(self.__orig_class__)
             if orig_class_args:
-                self._schema_annotations = orig_class_args[0]
+                self._schema_annotations: Type[_Schema] = orig_class_args[0]
                 validate_schema(
                     self._schema_annotations.get_structtype(),
                     deepcopy(self.schema),