Fix activation function in dpsgd conv model test.

google · Oct 28, 2024 · 521333e · 521333e
1 parent d2dc4ac
commit 521333e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tf_shell_ml/test/dpsgd_conv_model_local_test.py b/tf_shell_ml/test/dpsgd_conv_model_local_test.py
@@ -69,7 +69,8 @@ def _test_model(self, disable_encryption, disable_masking, disable_noise, cache)
                 tf_shell_ml.Flatten(),
                 tf_shell_ml.ShellDense(
                     16,
-                    activation=tf.nn.softmax,
+                    activation=tf_shell_ml.relu,
+                    activation_deriv=tf_shell_ml.relu_deriv,
                 ),
                 tf_shell_ml.ShellDense(
                     10,