# Transformer architecture not working on toy problem

My transformer is not working on a toy problem.

### Toy problem

Input : Sequence of random integer, one-hot-encoded. Example :

[[0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0]
]


Output : 1 if the first random integer is < 5, 0 otherwise. For the given example, the output would be 1.

### Problem

Running my Transformer on this toy problem works perfectly. However, running with a bigger size (one-hot encoding on 800 dimension) does not work anymore : the network always output similar results, no matter the input.

Note : I did change the condition of the output to be x < 800, in order to have class balance.

Why my transformer architecture is not working with this variation of the toy problem ?

### Another problem

I also tried to use 2 inputs (similarly generated, with size = 10) and generate the output based on these 2 inputs (1 if the first random integer of both input is < 5 or both > 5, 0 otherwise)