Solved by Orcun_Deniz in post #2 
              
              
                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)
        …