# Codeblock 6
class FullyConnected(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linear0 = nn.Linear(in_features=1024*7*7, out_features=4096)   #(1)
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.1)                  #(2)
        self.dropout = nn.Dropout(p=0.5)                                    #(3)
        self.linear1 = nn.Linear(in_features=4096, out_features=(C+B*5)*S*S)#(4)
    
    def forward(self, x):
        print(f'original\t: {x.size()}')
        
        x = self.linear0(x)
        print(f'after linear0\t: {x.size()}')
        
        x = self.leaky_relu(x)
        x = self.dropout(x)
        
        x = self.linear1(x)
        print(f'after linear1\t: {x.size()}')
        
        return x