Update models/vae.py

models/vae.py

@@ -0,0 +1,19 @@
+import torch
+import pytorch_lightning as pl
+
+class VAEModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        # Your VAE implementation here
+        self.encoder = torch.nn.Sequential(
+            torch.nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
+            torch.nn.ReLU()
+        )
+        self.decoder = torch.nn.Sequential(
+            torch.nn.ConvTranspose2d(16, 1, kernel_size=3, stride=2, padding=1),
+            torch.nn.Sigmoid()
+        )
+    
+    def forward(self, x):
+        z = self.encoder(x)
+        return self.decoder(z)