numpy array转pytorch dataloader

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1, test_size=0.3)
# X_train=ndarray(N,3,62,47), X_test=ndarray(N,3,62,47), y_train=ndarray(N,2), y_test=ndarray(N,2)

X_train, X_test = torch.FloatTensor(X_train), torch.FloatTensor(X_test)
y_train, y_test = torch.LongTensor(y_train), torch.LongTensor(y_test)
train_dataset = TensorDataset(X_train, y_train)
test_dataset = TensorDataset(X_test, y_test)

train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True)

for i, (images, labels) in enumerate(train_loader):
	pass