Thanks for the reply @ptrblck. It actually works with my dataset (dataSet from ImageFolder with transformations applied), but in the forward(X) method of the neural net, It passes a List of X, y tensors, instead of just X tensors, which is weird. I do Grid Search with Skorch, and here is the sample code:
start = time.time()
with torch.no_grad():

    net = NeuralNet(model, criterion= nn.CrossEntropyLoss,
                max_epochs = args.epochs,
                lr = args.learning_rate,
                batch_size = 32,
                optimizer=optim.SGD,
                optimizer_momentum = 0.09,
                iterator_train__shuffle = True,
                iterator_train__num_workers = 4,
                iterator_valid__shuffle = True,
                iterator_valid__num_workers = 4,
                train_split= predefined_split(valid_dataset),
                callbacks= [ lr_scheduler, epoch_acc, checkpoint],
                device = 'cuda')

#pipe = Pipeline([('scale',StandardScaler()),
#                 ('net',net)])
#pipe.fit()
params = {
        'lr':[0.01,0.02],
        'max_epochs':[10,20],
        'module_num_units': [10,20],
        }
gs = GridSearchCV(net, params,refit=False, cv=5, scoring='accuracy')

gs.fit(train_dataset, y=None)
end = time.time()
print("Total training time: "+ str(end - start))

Is there a specific way I need to specify to send just the corresponding X tensor in my forward function? Thanks