Solved this. I’ve been using flatten layer wrong by flattening through all dimensions. Changed the methods in model like;
    def _convs(self, image):
        image = image / 127.5 - 1
        conv1 = F.elu(self.conv_1(image), alpha=0.3)
        conv2 = F.elu(self.conv_2(conv1), alpha=0.3)
        conv3 = F.elu(self.conv_3(conv2), alpha=0.3)
        conv4 = F.elu(self.conv_4(conv3), alpha=0.3)
        conv5 = F.elu(self.conv_5(conv4), alpha=0.3)
        drop = self.drop(conv5)
        flat = torch.flatten(drop, start_dim=1, end_dim=3) ##Correction
        self.flattened_size = flat.shape[1] ##Correction

        return flat

    def forward(self, image):

        flat = self._convs(image)
        lin1 = F.elu(self.lin1(flat), alpha=0.3)
        lin2 = F.elu(self.lin2(lin1), alpha=0.3)
        lin3 = F.elu(self.lin3(lin2), alpha=0.3)
        lin4 = self.lin4(lin3)

        return lin4.squeeze() ## Minor Correction