Update testvae.py

testvae.py

@@ -11,18 +11,22 @@ import gc
 class TestModel(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.conv1 = torch.nn.Conv2d(3, 16, 5, 1, 2, bias=False)
+        self.start = torch.nn.Conv2d(3, 16, 3, 1, 1, bias=False)
+        self.conv1 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=False)
         self.conv2 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=False)
-        self.conv3 = torch.nn.Conv2d(16, 3, 3, 1, 1, bias=True)
+        self.conv3 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=False)
+        self.final = torch.nn.Conv2d(16, 3, 3, 1, 1, bias=False)
         self.bn1 = torch.nn.BatchNorm2d(16)
         self.bn2 = torch.nn.BatchNorm2d(16)
 
     def forward(self, x):
-        x = self.conv1(x)
+        x = self.start(x)
         x = self.bn1(x)
-        x = self.conv2(x)
+        x = self.conv1(x) + x
+        x = self.conv2(x) + x
+        x = self.conv3(x) + x
         x = self.bn2(x)
-        x = self.conv3(x)
+        x = self.final(x)
         x = torch.clamp(x, -1, 1)
         return x