Check the shape of the activation after flattening it and before passing it to self.fc1:
x = torch.flatten(x, 1)
print(x.shape)

This should give you a shape as [batch_size, num_features], where num_features should be equal to the input features of self.fc1.
Based on the error message you are seeing I guess x would have the shape [batch_size, 512], while you set in_features=256 for self.fc1, which causes the shape mismatch.
PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier.