Fix value fn

thatguy11325 · Mar 22, 2024 · 5ebb02d · 5ebb02d
1 parent 61226ef
commit 5ebb02d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pokemonred_puffer/policies/multi_convolutional.py b/pokemonred_puffer/policies/multi_convolutional.py
@@ -50,7 +50,7 @@ def __init__(
         )
 
         self.actor = nn.LazyLinear(self.num_actions)
-        self.value_fn = nn.LazyLinear(output_size, 1)
+        self.value_fn = nn.LazyLinear(1)
 
     def encode_observations(self, observations):
         observations = unpack_batched_obs(observations, self.unflatten_context)