Based on the error message the shape mismatch is created in the first linear layer:
torch.nn.Flatten(),
nn.Linear(64, 1024, bias = True),

so you could set the in_features to 1024 and rerun the script:
nn.Linear(1024, 1024, bias = True),