Update

[ghstack-poisoned]
pytorch · Nov 14, 2024 · 8fe1b84 · 8fe1b84
1 parent 819c77f
commit 8fe1b84
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
@@ -79,13 +79,11 @@ def __init__(
         self, type: InteractionType | str | None = InteractionType.DETERMINISTIC
     ) -> None:
         super().__init__()
-        if not isinstance(type, InteractionType):
+        if not isinstance(type, InteractionType) and type is not None:
             if isinstance(type, str):
                 type = InteractionType(type.lower())
             else:
-                raise TypeError(
-                    "Invalid type: only str or InteractionType are accepted."
-                )
+                raise ValueError(f"{type} is not a valid InteractionType")
         self.type = type
 
     def clone(self) -> set_interaction_type:

diff --git a/test/test_nn.py b/test/test_nn.py
@@ -88,9 +88,8 @@ def test_from_str_correct_conversion(self, str_and_expected_type):
 
     @pytest.mark.parametrize("unsupported_type_str", ["foo"])
     def test_from_str_correct_raise(self, unsupported_type_str):
-        with pytest.raises(ValueError) as err:
+        with pytest.raises(ValueError, match=" is not a valid InteractionType"):
             InteractionType.from_str(unsupported_type_str)
-        assert unsupported_type_str in str(err) and "is unsupported" in str(err)
 
 
 class TestTDModule: