output_dim is a hparam

matsengrp · Sep 27, 2024 · e0084dd · e0084dd
1 parent cbd8484
commit e0084dd
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/netam/models.py b/netam/models.py
@@ -598,6 +598,7 @@ def hyperparameters(self):
             "dim_feedforward": self.dim_feedforward,
             "layer_count": self.encoder.num_layers,
             "dropout_prob": self.pos_encoder.dropout.p,
+            "output_dim": self.linear.out_features,
         }
 
     def init_weights(self) -> None: