Fix test models to output logits and work with default loss functions

SeldonIO · Oct 12, 2023 · 6aca56d · 6aca56d
1 parent 7f1c2f3
commit 6aca56d
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/alibi/explainers/tests/test_simiarlity/conftest.py b/alibi/explainers/tests/test_simiarlity/conftest.py
@@ -151,7 +151,6 @@ def tf_linear_model(input_shape, output_shape):
         tf.keras.layers.InputLayer(input_shape=input_shape),
         tf.keras.layers.Flatten(),
         tf.keras.layers.Dense(output_shape),
-        tf.keras.layers.Softmax()
     ])
 
 
@@ -167,7 +166,6 @@ def __init__(self, input_shape, output_shape):
             self.linear_stack = nn.Sequential(
                 nn.Flatten(start_dim=1),
                 nn.Linear(input_shape, output_shape),
-                nn.Softmax()
             )
 
         def forward(self, x):