From baf6d9e21d1f11b8e0fbed7ac9c5e880bba70c3a Mon Sep 17 00:00:00 2001
From: David Emerson <43939939+emersodb@users.noreply.github.com>
Date: Tue, 3 Jun 2025 22:32:29 -0400
Subject: [PATCH 1/2] Fixing typing annotation in Bernoulli sample

---
 fl4health/utils/functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fl4health/utils/functions.py b/fl4health/utils/functions.py
index b0826d2fb..97f13afbe 100644
--- a/fl4health/utils/functions.py
+++ b/fl4health/utils/functions.py
@@ -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

From 4f25c070ccb63ed04773b53a6cd10d07865cb43a Mon Sep 17 00:00:00 2001
From: David Emerson <43939939+emersodb@users.noreply.github.com>
Date: Wed, 4 Jun 2025 10:56:24 -0400
Subject: [PATCH 2/2] Removing the type ignores

---
 fl4health/utils/functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fl4health/utils/functions.py b/fl4health/utils/functions.py
index 97f13afbe..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