python - 无法让 skorch 在 GPU 上运行

Question

我的 skorch 神经网络 GridsearchCV 代码正在运行，但它似乎不在 GPU 上，所以速度很慢。

任何帮助表示赞赏。

每个 epoch 需要 13 秒。我知道 pytorch 在 GPU 上运行得更快。

代码：

import skorch
import torch
import torch.nn as nn

from skorch.classifier import NeuralNetBinaryClassifier
from sklearn import model_selection
from torch.nn import functional as F

cuda = torch.device('cuda')
parameters = {'lr':[1e-4,1e-5]}

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.lin1 = nn.Linear(50, 30)
        self.lin2 = nn.Linear(30, 1)

    def forward(self, x):
        x = F.tanh(self.lin1(x))
        return self.lin2(x).squeeze()

class NLL(torch.nn.modules.loss._Loss):
    def forward(self, input, target):
        loss = F.nll_loss(input.unsqueeze(1),target)
        return loss   


model = Model()
model.to(cuda)

cl1 = NeuralNetBinaryClassifier(
    module=model,
    criterion=NLL,
    optimizer=torch.optim.Adam,
    max_epochs=100,
    train_split=None,
    device='cuda',
    )  

clf = model_selection.GridSearchCV(cl1, parameters, 
    iid=False, 
    scoring='neg_log_loss',
    cv=ps,
    verbose=1,
    n_jobs=1,
    return_train_score=True,
    )

clf.fit(X, Y)