The size of tensor a (4) must match the size of tensor b (128) at non-singleton dimension 0

When evaluation loop is applied over the trainset I’m getting this error

Hi, can you share your notebook with a public access link?