diff --git a/src/liecasadi/so3.py b/src/liecasadi/so3.py
index 9246631..ff3f457 100644
--- a/src/liecasadi/so3.py
+++ b/src/liecasadi/so3.py
@@ -4,6 +4,7 @@
 
 import dataclasses
 from dataclasses import field
+from typing import Union
 
 import casadi as cs
 import numpy as np
@@ -113,7 +114,7 @@ def quaternion_derivative(
         self,
         omega: Vector,
         omega_in_body_fixed: bool = False,
-        baumgarte_coefficient: float = None,
+        baumgarte_coefficient: Union[float, None] = None,
     ):
 
         if baumgarte_coefficient is not None: