diff --git a/pokemonred_puffer/policies/multi_convolutional.py b/pokemonred_puffer/policies/multi_convolutional.py
index 423e86a..fdc2e75 100644
--- 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)