Add __call__ typing to CategoricalPolicy

takuseno · Aug 6, 2023 · fc739c7 · fc739c7
1 parent bf1632f
commit fc739c7
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/d3rlpy/models/torch/policies.py b/d3rlpy/models/torch/policies.py
@@ -146,3 +146,6 @@ def __init__(self, encoder: Encoder, hidden_size: int, action_size: int):
 
     def forward(self, x: torch.Tensor) -> Categorical:
         return Categorical(logits=self._fc(self._encoder(x)))
+
+    def __call__(self, x: torch.Tensor) -> Categorical:
+        return super().__call__(x)