Models should not compute loss when encrypted backprop enabled.

google · Oct 30, 2024 · a1b9afa · a1b9afa
1 parent da280db
commit a1b9afa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tf_shell_ml/dpsgd_sequential_model.py b/tf_shell_ml/dpsgd_sequential_model.py
@@ -247,10 +247,10 @@ def rebalance(x, s):
         for metric in self.metrics:
             if metric.name == "loss":
                 if self.disable_encryption:
-                    metric.update_state(0)  # Loss is not known when encrypted.
-                else:
                     loss = self.loss_fn(y, y_pred)
                     metric.update_state(loss)
+                else:
+                    metric.update_state(0)  # Loss is not known when encrypted.
             else:
                 if self.disable_encryption:
                     metric.update_state(y, y_pred)

diff --git a/tf_shell_ml/postscale_sequential_model.py b/tf_shell_ml/postscale_sequential_model.py
@@ -249,10 +249,10 @@ def rebalance(x, s):
         for metric in self.metrics:
             if metric.name == "loss":
                 if self.disable_encryption:
-                    metric.update_state(0)  # Loss is not known when encrypted.
-                else:
                     loss = self.loss_fn(y, y_pred)
                     metric.update_state(loss)
+                else:
+                    metric.update_state(0)  # Loss is not known when encrypted.
             else:
                 if self.disable_encryption:
                     metric.update_state(y, y_pred)