better error message

bjtuxck · Oct 10, 2018 · 4212a06 · 4212a06
1 parent d07d16a
commit 4212a06
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/cleverhans/model.py b/cleverhans/model.py
@@ -45,7 +45,12 @@ def get_logits(self, x, **kwargs):
     :return: A symbolic representation (Tensor) of the output logits
     (i.e., the values fed as inputs to the softmax layer).
     """
-    return self.fprop(x, **kwargs)[self.O_LOGITS]
+    outputs = self.fprop(x, **kwargs)
+    if self.O_LOGITS in outputs:
+      return outputs[self.O_LOGITS]
+    raise NotImplementedError(str(type(self)) + "must implement `get_logits`"
+                              " or must define a " + self.O_LOGITS +
+                              " output in `fprop`")
 
   def get_probs(self, x, **kwargs):
     """