[RLlib] Fix configs.py for framework != torch (e.g. tf2). (#35975)

ray-project · Sep 15, 2023 · 49cd01f · 49cd01f
1 parent 515fbe8
commit 49cd01f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/rllib/core/models/configs.py b/rllib/core/models/configs.py
@@ -520,7 +520,7 @@ def _validate(self, framework: str = "torch"):
 
     @_framework_implemented()
     def build(self, framework: str = "torch") -> "Model":
-        self._validate()
+        self._validate(framework)
 
         if framework == "torch":
             from ray.rllib.core.models.torch.heads import TorchCNNTransposeHead
@@ -682,7 +682,7 @@ def _validate(self, framework: str = "torch"):
 
     @_framework_implemented()
     def build(self, framework: str = "torch") -> "Model":
-        self._validate()
+        self._validate(framework)
 
         if framework == "torch":
             from ray.rllib.core.models.torch.encoder import TorchCNNEncoder
@@ -746,7 +746,7 @@ class MLPEncoderConfig(_MLPConfig):
 
     @_framework_implemented()
     def build(self, framework: str = "torch") -> "Encoder":
-        self._validate()
+        self._validate(framework)
 
         if framework == "torch":
             from ray.rllib.core.models.torch.encoder import TorchMLPEncoder