# How to Connect Convolutional layer to Fully Connected layer in Pytorch while Implementing SRGAN

I was implementing the SRGAN in PyTorch but while implementing the discriminator I was confused about how to add a fully connected layer of 1024 units after the final convolutional layer My input data shape:(1,3,256,256)

After passing this data through the conv layers I get a data shape: torch.Size([1, 512, 16, 16])

Code:

class Discriminator(nn.Module):

def __init__(self):
super(Discriminator,self).__init__()
self.sm = nn.Sigmoid()
self.net = nn.Sequential(
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),

nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),

nn.Linear(<ADD AN INPUT SHAPE HERE>,1024),
nn.LeakyReLU(0.2),
nn.Linear(1024,1)
)

def forward(self,x):
x = self.sm(self.net(x))
x = (x)
return x


You simply reshape the tensor to (batch_size, n_nodes) using tensor.view(). In your specific case this would be x.view(x.size()[0], -1).