the problem is on my CrossEntropy. Here is an example of my logits and target I pass to CrossEntropy
tensor([[[-0.0006, -0.0089],
         [ 0.0044, -0.0024],
         [-0.0049,  0.0009],
         [-0.0042, -0.0018],
         [-0.0041, -0.0021],
         [ 0.0028, -0.0017],
         [-0.0033, -0.0095],
         [-0.0022, -0.0034],
         [-0.0064, -0.0046],
         [ 0.0048, -0.0077],
         [-0.0050, -0.0023],
         [ 0.0005, -0.0038],
         [-0.0037, -0.0002],
         [ 0.0008, -0.0033],
         [-0.0006, -0.0026],
         [-0.0013, -0.0122],
         [-0.0124, -0.0082],
         [ 0.0015, -0.0039],
         [ 0.0065, -0.0008],
         [ 0.0003, -0.0117],
         [ 0.0024, -0.0064],
         [-0.0017, -0.0041],
         [-0.0045, -0.0033],
         [ 0.0049,  0.0017],
         [-0.0013, -0.0082],
         [-0.0001, -0.0014]]], device='cuda:1', grad_fn=<AddBackward0>)
tensor([0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0,
        0, 1], device='cuda:1')