From be1c8cd235dc27de8d86657af7dc97e66cd7f923 Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Mon, 29 Apr 2024 13:37:26 +0200 Subject: [PATCH] DQN: * Fix input validation * Fix output_dim not being set if features_only=True and output_dim_added_layer not None --- examples/atari/atari_network.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/atari/atari_network.py b/examples/atari/atari_network.py index ea900e9..4f2a560 100644 --- a/examples/atari/atari_network.py +++ b/examples/atari/atari_network.py @@ -66,7 +66,7 @@ class DQN(NetBase[Any]): layer_init: Callable[[nn.Module], nn.Module] = lambda x: x, ) -> None: # TODO: Add docstring - if features_only and output_dim_added_layer is not None: + if not features_only and output_dim_added_layer is not None: raise ValueError( "Should not provide explicit output dimension using `output_dim_added_layer` when `features_only` is true.", ) @@ -98,6 +98,7 @@ class DQN(NetBase[Any]): layer_init(nn.Linear(base_cnn_output_dim, output_dim_added_layer)), nn.ReLU(inplace=True), ) + self.output_dim = output_dim_added_layer else: self.output_dim = base_cnn_output_dim