added a FFN in CNN

config.yml

@@ -57,6 +57,7 @@ BidirectionalLSTM:
 RNN:
   HiddenSize: 35
   NumLayers: 4
+  NumFFN: 8
   Initialization: None
 
 #Name of directory containing logs

model.py

@@ -42,16 +42,34 @@ class RNN(nn.Module):
         super(RNN, self).__init__()
         self.hidden_size = cfg["RNN"]["HiddenSize"]
         self.num_layers = cfg["RNN"]["NumLayers"]
+        self.num_ffn = cfg["RNN"]["NumFFN"]
         # RNN
         self.rnn = nn.Sequential(
             nn.RNN(input_size, self.hidden_size, self.num_layers, batch_first=True, nonlinearity='relu'),
             nn.Dropout(p=0.2)
         )
+
+
+        self.fc = nn.Sequential()
         
-        self.fc = nn.Sequential(
-            nn.Linear(self.hidden_size, 1),
-            nn.Dropout(p=0.2)
+        for layer in range(self.num_ffn):
+            self.fc.add_module(
+                f"linear_{layer}", nn.Linear(self.hidden_size, self.hidden_size)
+            )
+            self.ffn.add_module(
+                f"relu_{layer}", 
+                nn.ReLU()
+            )
+            self.ffn.add_module(
+                f"dropout_{layer}", 
+                nn.Dropout(p=0.2)
+            )
+        
+        self.fc.add_module(
+            "last_linear",
+            nn.Linear(self.hidden_size, 1)
         )
+
     
     def forward(self, x):
         use_cuda = torch.cuda.is_available()
@@ -61,8 +79,7 @@ class RNN(nn.Module):
             device = toch.device('cpu')
         # Initialize hidden state with zeros
         h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
-            
-        # One time step
+    
         out, hn = self.rnn(x, h0)
         out = self.fc(out)
         return out