Matrix incompatible as CNN input

Question

I'm using now pytorch for image classification from numpy file, I have 2 separated files one for the data and the other for labels, so I applied a split by: sklearn.model_selection

X_train, X_test, y_train, y_test = train_test_split(data, labels, test_size=0.2, shuffle=True, stratify=labels )

The shape of results are:

X_train.shape   torch.Size([320, 224, 224, 3])     
X_test.shape   torch.Size([80, 224, 224, 3])
y_train.shape  torch.Size([320])
y_test.shape   torch.Size([80])

I reshape X_train and X_test.shape to ([320, 3, 224, 224]) , ([80, 3,224, 224]) Now I want to put these data into AlexNet model (for example) and apply the test to get the accuracy, but I didn't find a good solution!!! Help me please...

Now I want to put these data into AlexNet model (for example) and apply the test to get the accuracy, but I didn't find a good solution!!! Help me please...