diff --git a/test/test_ops.py b/test/test_ops.py index f216608b19..cea8081777 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -59,6 +59,8 @@ class TestOps(unittest.TestCase): helper_test_op([(45,65), (65,100)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-5, gpu=self.gpu) def test_sum(self): helper_test_op([(45,1)], lambda x: x.sum(), Tensor.sum, atol=1e-5, gpu=self.gpu) + def test_logsoftmax(self): + helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.logsoftmax, atol=1e-5, gpu=self.gpu) def test_conv2d(self): for bs in [1,8]: