diff --git a/fl4health/utils/functions.py b/fl4health/utils/functions.py
index b0826d2fb..835fc47f4 100644
--- a/fl4health/utils/functions.py
+++ b/fl4health/utils/functions.py
@@ -19,7 +19,7 @@ class BernoulliSample(torch.autograd.Function):
     """
 
     @staticmethod
-    def forward(bernoulli_probs: torch.Tensor) -> torch.Tensor:  # type: ignore
+    def forward(bernoulli_probs: torch.Tensor) -> torch.Tensor:
         return torch.bernoulli(input=bernoulli_probs)
 
     @staticmethod
@@ -33,7 +33,7 @@ def setup_context(ctx: Any, inputs: tuple[torch.Tensor], output: torch.Tensor) -
     # This method determines the "gradient" of the BernoulliSample function.
     # grad_output is supposed to be the gradient w.r.t. the output of the forward method.
     @staticmethod
-    def backward(ctx: torch.Any, grad_output: torch.Tensor) -> torch.Tensor:  # type: ignore
+    def backward(ctx: Any, grad_output: torch.Tensor) -> torch.Tensor:  # type: ignore
         # ctx.saved_tensors is a tuple (of length 1 in this case). Hence the indexing here.
         bernoulli_probs = ctx.saved_tensors[0]
         return bernoulli_probs * grad_output