Fix tests

Author: araffin

tests/test_distributions.py

@@ -82,19 +82,19 @@ def test_distribution_must_be_initialized(self):
 
         DIMS = 2
         dist = MaskableCategoricalDistribution(DIMS)
-        with pytest.raises(AssertionError):
+        with pytest.raises(AttributeError):
             dist.log_prob(th.randint(DIMS - 1, (1, 3)))
 
-        with pytest.raises(AssertionError):
+        with pytest.raises(AttributeError):
             dist.entropy()
 
-        with pytest.raises(AssertionError):
+        with pytest.raises(AttributeError):
             dist.sample()
 
-        with pytest.raises(AssertionError):
+        with pytest.raises(AttributeError):
             dist.mode()
 
-        with pytest.raises(AssertionError):
+        with pytest.raises(AttributeError):
             dist.apply_masking(None)
 
         # But now we can