Thanks for the reply. How would you suggest me to edit the final layer? What should the output dimension be?
Since the number of rows of my input data is 10000, I am not sure how the output layer does not include that number.
Here are the dimensions after each layer:
Dimension of input: torch.Size([10, 10000])
Dimension after embedding: torch.Size([10, 10000, 150])
Dimension after layer1: torch.Size([10, 10000, 512])
Dimension after layer2: torch.Size([10, 10000, 512])
Dimension after final layer: torch.Size([10, 10000, 2])