You could write a basic custom dataset and use that with a dataloader.
class DS(torch.utils.data.Dataset):
    def __init__(this, X=None, y=None, mode="train"):
        this.mode = mode
        this.X = X #Maybe do reshaping here
        if mode == "train":
            this.y = y

    def __len__(this):
        return this.X.shape[0]

    def __getitem__(this, idx):
        if this.mode == "train":
            return torch.FloatTensor(this.X[idx]), torch.LongTensor(this.y[idx]) #or torch.FloatTensor(this.y[idx]) depending on use case
        else:
            return torch.FloatTensor(this.X[idx])

tr_data_setup = DS(X_train, y_train.reshape(-1,1)) #
trainloader = torch.data.utils.DataLoader(tr_data_setup, batch_size=100, ......)

You could also expand this to perform augmentations on the image if necessary.