generator.py

class Generator(nn.Module):
    def __init__(self, random_size, hidden_size, image_size):
        super(Generator, self).__init__()
        #random _size -> 64, 
        #Input is random noise with a fixed size
        self.fc1 = nn.Linear(random_size, hidden_size) 
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        #Final output is of shape equal to image size, 
        #For MNIST -> 784 (as flatten)
        self.fc3 = nn.Linear(hidden_size, image_size) 
                                                      
        self.tanh = nn.Tanh()
    def forward(self, x):
        out_fc1 = self.fc1(x)
        out_relu = self.relu(out_fc1)
        out_fc2 = self.fc2(out_relu)
        out_relu2 = self.relu(out_fc2)
        out_fc3 = self.fc3(out_relu2)
        out = self.tanh(out_fc3)
        return(out)