Update pipeline.py

pipeline.py

@@ -1,9 +1,29 @@
 import torch
 from PIL import Image
 from torchvision import transforms
+import torch.nn as nn
 
 class Generator(nn.Module):
-    # The Generator class remains the same as provided
+    def __init__(self, input_size, output_channels):
+        super(Generator, self).__init__()
+        
+        # Define the architecture of the generator
+        self.model = nn.Sequential(
+            nn.Linear(input_size, 128),         # Input layer
+            nn.LeakyReLU(0.2),                 # Activation function
+            nn.Linear(128, 256),               # Hidden layer
+            nn.BatchNorm1d(256),               # Batch normalization
+            nn.LeakyReLU(0.2),                 # Activation function
+            nn.Linear(256, 512),               # Hidden layer
+            nn.BatchNorm1d(512),               # Batch normalization
+            nn.LeakyReLU(0.2),                 # Activation function
+            nn.Linear(512, output_channels),    # Output layer
+            nn.Tanh()                          # Tanh activation for output
+        )
+
+    def forward(self, x):
+        # Forward pass through the generator network
+        return self.model(x)
 
 class PreTrainedPipeline():
     def __init__(self, path=""):