(#7) Attack: Support for non-layer-based models

betarixm · May 16, 2022 · 3c9238a · 3c9238a
1 parent 17508b0
commit 3c9238a
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/typings/models.py b/src/typings/models.py
@@ -31,8 +31,13 @@ def __init__(
 
         self.__input_shape = input_shape
 
-        self.__model: keras.Model = keras.Sequential(
-            [keras.Input(self.__input_shape), *self._model().layers], name=self._name
+        self.__model: keras.Model = (
+            keras.Sequential(
+                [keras.Input(self.__input_shape), *self._model().layers],
+                name=self._name,
+            )
+            if len(self._model().layers) != 0
+            else self._model()
         )
 
         self.intensity: float = intensity
@@ -115,7 +120,7 @@ def compile(self):
         )
 
     def predict(self, inputs):
-        outs = self.__model.predict(inputs)
+        outs = self.__model(inputs)
         return inputs + (outs - inputs) * self.intensity if self.intensity < 1 else outs
 
     def train(self, epochs: int = 100):