📖 Model Summarization

LuluW8071 · Mar 21, 2024 · bb96380 · bb96380
1 parent 43d1cb5
commit bb96380
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/neuralnet/model.py b/neuralnet/model.py
@@ -100,12 +100,12 @@ def _init_hidden(self, batch_size):
                 torch.zeros(n*1, batch_size, hs))
 
     def forward(self, x, hidden):
-        x = x.squeeze(1)  # batch, feature, time
-        x = self.cnn(x)  # batch, time, feature
-        x = self.dense(x)  # batch, time, feature
-        x = x.transpose(0, 1)  # time, batch, feature
+        x = x.squeeze(1)                                    # batch, feature, time
+        x = self.cnn(x)                                     # batch, time, feature
+        x = self.dense(x)                                   # batch, time, feature
+        x = x.transpose(0, 1)                               # time, batch, feature
         out, (hn, cn) = self.lstm(x, hidden)
-        x = self.dropout2(F.gelu(self.layer_norm2(out)))  # (time, batch, n_class)
+        x = self.dropout2(F.gelu(self.layer_norm2(out)))    # (time, batch, n_class)
         return self.final_fc(x), (hn, cn)
 
     # Model Summarization