[BUG] Cannot process using GPU

Describe the bug I have trained the model NBEATS for a week, things worked properly if I train the model on single run. However, when I need to do gridsearch on this model, Data have just loaded on GPU, but calculating on CPU only, so it is very slow. I run the code on Jupyter lab.

To Reproduce I trained data with normal random variable with some available correlation, and scaled using sklearn.StandardScaler()

train_y, val_y = train_test_split(y, test_size = 0.3)
train_x, val_x = train_test_split(X, test_size = 0.3)

train_y.shape
>>> (1400,1)
val_y.shape
>>>(600,1)
train_x.shape
>>> (1400,22)
val_x.shape
>>>(600,22)

params2 = {
  'input_chunk_length':[12,18],
  'output_chunk_length': [3],
  'generic_architecture': [True],
  'num_stacks': [2,4,6],
  'num_blocks' : [2,4,6],
  'num_layers' : [2,4,6], 
  'layer_widths': [22],
  'batch_size': [128],
  'random_state' : [0]
}

gs_test2 = NBEATSModel.gridsearch(
     params2, 
     series = train_y, 
     past_covariates = train_x, 
     verbose = True, 
     forecast_horizon = 12
)

Expected behavior Expected both loading and training data on GPU

System (please complete the following information):

• Python version: 3.8.11
• darts version: 0.12.0
• cudatoolkit: 11.1.74
• jupyterlab: 3.1.7
• numpy: 1.19.5
• pytorch: 1.8.2
• scikit-learn: 0.23.2 Hard drive system
• CPU: i7-8700
• GPU: RTX-2070

Thank for your help!

[hrzn]

You have to include the GPU config in the hyper-parameters grid of the gridsearch. Something like this:

hyper_params = {
  ...,
  'pl_trainer_kwargs': [{
      "accelerator": "gpu",
      "gpus": [0]
    }],
}

... = ModelClass.gridsearch(parameters=hyper_params, ...)

[smhoma]

You have to include the GPU config in the hyper-parameters grid of the gridsearch. Something like this:

hyper_params = {
  ...,
  'pl_trainer_kwargs': [{
      "accelerator": "gpu",
      "gpus": [0]
    }],
}

... = ModelClass.gridsearch(parameters=hyper_params, ...)

It worked. thank you very much.