fail fast if unproper strategy passed in (#881)

Summary:
Pull Request resolved: #881

tsia

Reviewed By: diego-urgell

Differential Revision: D61241468

fbshipit-source-id: 05e707ca1307961668a5b674e101e9b2165944d1

Author: JKSenthil

tests/utils/test_prepare_module.py

@@ -75,6 +75,15 @@ def test_prepare_module_strategy_invalid_str(self) -> None:
                 strategy="foo",
             )
 
+    def test_prepare_module_invalid_strategy(self) -> None:
+        with self.assertRaisesRegex(ValueError, "Unknown strategy received"):
+            prepare_module(
+                module=torch.nn.Linear(2, 2),
+                device=init_from_env(),
+                # pyre-ignore: Incompatible parameter type [6] (intentional to test error raised)
+                strategy={"_strategy_": "DDPStrategy"},
+            )
+
     def test_prepare_noop(self) -> None:
         device = torch.device("cuda")  # Suppose init_from_env returns cuda
 

torchtnt/utils/prepare_module.py

@@ -312,6 +312,11 @@ def prepare_module(
     """
 
     if strategy:
+        if not isinstance(strategy, str) and not isinstance(strategy, Strategy):
+            raise ValueError(
+                f"Unknown strategy received: {strategy}. Expect either str or Strategy dataclass"
+            )
+
         if isinstance(strategy, str):
             strategy = convert_str_to_strategy(strategy)
         if isinstance(strategy, DDPStrategy):