Here is where i defined predictions (it is a multi-class classifier):
def test(dataset, dataloader):

    # switch to test mode
    net.eval()  

    # initialize predictions
    predictions = torch.zeros(len(dataset), dtype=torch.int64)
    sample_counter = 0

    # do not accumulate gradients (faster)
    with torch.no_grad():

        # test all batches
        for batch in dataloader:

            # get data from dataloader [ignore labels/targets as they are not used in test mode]
            inputs = batch[0]

            # move data to device
            inputs = inputs.to(device, non_blocking=True)

            # forward pass
            outputs = net(inputs)

            # store predictions
            outputs_max = torch.argmax(outputs, dim=1)
            for output in outputs_max:
                predictions[sample_counter] = output
                sample_counter += 1

    return predictions

I tought it was a problem with random_split becouse this problem has rised since I used that function.
Previosly I was using the original test set from MNIST with the same code and had no issue.
Thanks for helping.