tensor([-2.5495, -2.6069, -2.0307, -1.6827, -2.8096, -2.8205, -2.7063, -1.8902,
        -2.6554, -2.0948], grad_fn=<SelectBackward0>) torch.Size([64, 10])