From 52c194ad8943e1abbaeca8ff87e5b0a31625bcc8 Mon Sep 17 00:00:00 2001 From: Gao Enhao Date: Fri, 17 Nov 2023 21:16:43 +0800 Subject: [PATCH] [MNT] add test_predict_proba --- tests/test_basic_nn.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_basic_nn.py b/tests/test_basic_nn.py index f344700..910a3bf 100644 --- a/tests/test_basic_nn.py +++ b/tests/test_basic_nn.py @@ -35,6 +35,13 @@ class TestBasicNN(object): assert len(predictions) == len(X) assert numpy.isin(predictions, list(range(10))).all() + # Test predict_proba method + def test_predict_proba(self, basic_nn_instance): + X = list(torch.randn(32, 1, 28, 28)) + predict_proba = basic_nn_instance.predict_proba(X=X) + assert len(predict_proba) == len(X) + assert ((0 <= predict_proba) & (predict_proba <= 1)).all() + # Test score method def test_score(self, basic_nn_instance): X = torch.randn(32, 1, 28, 28)