refactor annotation

Author: mkumar73

mambular/arch_utils/layer_utils/batch_ensemble_layer.py

@@ -167,7 +167,7 @@ def reset_parameters(self, scaling_init: Literal["ones", "random-signs", "normal
         if self.bias is not None:
             nn.init.zeros_(self.bias)
 
-    def forward(self, x: torch.Tensor, hidden: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, hidden: torch.Tensor = None) -> torch.Tensor:  # type: ignore
         """Forward pass for the BatchEnsembleRNNLayer.
 
         Parameters

mambular/arch_utils/transformer_utils.py

@@ -168,7 +168,7 @@ def __init__(
         else:
             raise ValueError(f"Invalid activation '{activation}'. Choose from 'relu' or 'gelu'.")
 
-    def forward(self, src, src_mask: torch.Tensor | None = None):
+    def forward(self, src, src_mask: torch.Tensor = None):  # type: ignore
         """Pass the input through the encoder layer.
 
         Parameters
@@ -313,7 +313,7 @@ def __init__(
 
         self.ensemble_projections = ensemble_projections
 
-    def forward(self, x, mask: torch.Tensor | None = None):
+    def forward(self, x, mask: torch.Tensor = None):  # type: ignore
         """Pass the input through the encoder layers in turn.
 
         Parameters