From e0084dd4f616f80c99e91b2301677da46d21968d Mon Sep 17 00:00:00 2001
From: Erick Matsen <ematsen@gmail.com>
Date: Fri, 27 Sep 2024 09:36:04 -0700
Subject: [PATCH] output_dim is a hparam

---
 netam/models.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/netam/models.py b/netam/models.py
index 5267dec8..17008795 100644
--- 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: