Is it working with Keras

Gurobi · Dec 4, 2024 · 172c02a · 172c02a
1 parent 46531fb
commit 172c02a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/gurobi_ml/keras/keras.py b/src/gurobi_ml/keras/keras.py
@@ -110,7 +110,7 @@ def _mip_model(self, **kwargs):
                 output = self._output
             if isinstance(step, keras.layers.InputLayer):
                 pass
-            elif isinstance(step, keras.layers.ReLU):
+            elif isinstance(step, (keras.layers.ReLU, keras.layers.Softmax)):
                 layer = self._add_activation_layer(
                     _input, self.act_dict["relu"](), output, name=f"relu{i}", **kwargs
                 )