diff --git a/bitsandbytes/research/autograd/_functions.py b/bitsandbytes/research/autograd/_functions.py
index 390f54fed..0a4b2a0fd 100644
--- a/bitsandbytes/research/autograd/_functions.py
+++ b/bitsandbytes/research/autograd/_functions.py
@@ -186,7 +186,7 @@ class SwitchBackBnb(torch.autograd.Function):
     @staticmethod
     # TODO: the B008 on the line below is a likely bug; the current implementation will
     #       have each SwitchBackBnb instance share a single MatmulLtState instance!!!
-    def forward(ctx, A, B, out=None, bias=None, state: MatmulLtState = None):
+    def forward(ctx, A, B, out=None, bias=None, state: Optional[MatmulLtState] = None):
         state = state or MatmulLtState()
 
         # default to pytorch behavior if inputs are empty