I changed it to this:
model = nn.Sequential(nn.Conv2d(3, 32, kernel_size=(3,3), stride=1),
                      nn.Conv2d(32, 32, kernel_size=(3,3), stride=1),
                      nn.MaxPool2d(kernel_size=(2,2), stride=2),
                      nn.ReLU(),
                      nn.Conv2d(32, 64, kernel_size=(3,3), stride=1),
                      nn.Conv2d(64, 64, kernel_size=(3,3), stride=1),
                      nn.MaxPool2d(kernel_size=(2,2), stride=2),
                      nn.ReLU(),
                      nn.Conv2d(64, 128, kernel_size=(3,3), stride=1),
                      nn.Conv2d(128, 128, kernel_size=(3,3), stride=1),
                      nn.Conv2d(128, 128, kernel_size=(3,3), stride=1),
                      nn.MaxPool2d(kernel_size=(2,2), stride=2),
                      nn.ReLU(),
                      nn.Conv2d(128, 256, kernel_size=(3,3), stride=1),
                      nn.Conv2d(256, 256, kernel_size=(3,3), stride=1),
                      nn.Conv2d(256, 256, kernel_size=(3,3), stride=1),
                      nn.MaxPool2d(kernel_size=(2,2), stride=2),
                      nn.ReLU(),
                      nn.Flatten(),
                      nn.Linear(256,256),
                      nn.Linear(256,256),
                      nn.ReLU(),
                      nn.Linear(256,2),
                      nn.LogSoftmax(dim=1))

But still getting this error:
RuntimeError: size mismatch, m1: [32 x 25600], m2: [256 x 256] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:290