* Fix input validation
  * Fix output_dim not being set if features_only=True and output_dim_added_layer not None
This commit is contained in:
Dominik Jain 2024-04-29 13:37:26 +02:00
parent 081adedc32
commit be1c8cd235

View File

@ -66,7 +66,7 @@ class DQN(NetBase[Any]):
layer_init: Callable[[nn.Module], nn.Module] = lambda x: x, layer_init: Callable[[nn.Module], nn.Module] = lambda x: x,
) -> None: ) -> None:
# TODO: Add docstring # 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( raise ValueError(
"Should not provide explicit output dimension using `output_dim_added_layer` when `features_only` is true.", "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)), layer_init(nn.Linear(base_cnn_output_dim, output_dim_added_layer)),
nn.ReLU(inplace=True), nn.ReLU(inplace=True),
) )
self.output_dim = output_dim_added_layer
else: else:
self.output_dim = base_cnn_output_dim self.output_dim = base_cnn_output_dim