CallableModelWrapper can have different layer names

bjtuxck · Jul 12, 2017 · 9abc88f · 9abc88f
1 parent 24100cb
commit 9abc88f
Showing 1 changed file with 0 additions and 4 deletions.
diff --git a/cleverhans/model.py b/cleverhans/model.py
@@ -79,10 +79,6 @@ def __init__(self, callable_fn, output_layer):
         :param output_layer: A string of the output layer returned by the
                              function. (Usually either "probs" or "logits".)
         """
-        if output_layer not in ['probs', 'logits']:
-            error_str = 'output_layer argument of CallableModelWrapper' \
-                        'should be the string "probs" or "logits".'
-            raise ValueError(error_str)
 
         self.output_layer = output_layer
         self.callable_fn = callable_fn