add cvae training

models/base.py

@@ -6,22 +6,24 @@ from typing import Any, Callable, List, Optional, Sequence
 
 class VariationalEncoder(nn.Module):
     
-    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
+    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
         super(VariationalEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
 
-        self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
-        self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.mean_branch = nn.Sequential(*layers).to(device)
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.var_branch = nn.Sequential(*layers).to(device)
 
         self.flatten_mean = nn.Flatten()
         self.flatten_var = nn.Flatten()
 
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         x = inputs
 
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
         mean = self.mean_branch(x)
         log_var = self.var_branch(x)