Electra - Cannot reshape a tensor with 1048576 elements to shape [1,512,16,64] (524288 elements)

[Marius Johan Schlichtkrull]

Cannot reshape a tensor with 1048576 elements to shape [1,512,16,64] (524288 elements)

When I try running the code I get this error, and I get the error because one of the values is half as big as it is supposed to.
1048576 / 524288 = 2 and 1 x 512 x 16 x 64 = 524288.
[batch_size=1, seq_length=512, num_attention_heads=16, width=64]

Therefore on of the shapes should be doubled, but I don't know which one.
All of this is happining in the transpose_for_scores function.

Edit the input tensor has a shape of [1024 x 1024], I don't know that is supposed to be less

[Marius Johan Schlichtkrull]

Update I think the error is because the from_tensor and to_tensor has a shape of [1024 x 1024] instead of [512 x 1024]