Remove test case too
@@ -59,5 +59,15 @@ class TestModelClass(unittest.TestCase):
self.assertTrue(context.exception)
+class TestCallableModelWrapperInitArguments(unittest.TestCase):
+ def test_output_layer(self):
+ def model():
+ return True
+
+ # The following two calls should not raise Exceptions
+ wrap = CallableModelWrapper(model, 'probs')
+ wrap = CallableModelWrapper(model, 'logits')
+
+
if __name__ == '__main__':
unittest.main()
