Выпуск обучающей модели pytorch на gpu

Question

Я пытаюсь реализовать дискриминатор основы c MNIST GAN в PyTorch. Когда я запускаю тренировку на процессоре, он работает без каких-либо проблем, давая желаемый результат. Тем не менее, когда я запускаю его на GPU, он показывает ошибку во время выполнения. Я вставляю код для моей модели, а также мои тренировки ниже, а также изменения, которые я сделал, чтобы попытаться запустить тренировку на GPU.

dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

def preprocess(x):
    return x.view(-1,1,28,28).to(dev)

discriminator = nn.Sequential(
    Lambda(preprocess),
    nn.Conv2d(1,64,3,stride=2,padding=1),
    nn.LeakyReLU(negative_slope=0.2),
    nn.Dropout(0.4),
    nn.Conv2d(64,64,3,stride=2,padding=1),
    nn.LeakyReLU(negative_slope=0.2),
    nn.Dropout(0.4),
    Lambda(lambda x:x.view(x.size(0),-1)),
    nn.Linear(3136,1),
    nn.Sigmoid()
)
loss = nn.BCELoss()
opt = optim.Adam(discriminator.parameters(),lr = 0.002)

discriminator.to(dev)

def train_discriminator(model, dataset,opt, n_iter=100, n_batch=256):
    half_batch = int(n_batch / 2)
    for i in range(n_iter):
        X_real, y_real = generate_real_samples(dataset, half_batch)
        error = loss(model(X_real),y_real)
        error.backward()
        X_fake, y_fake = generate_fake_samples(half_batch)
        error = loss(model(X_fake),y_fake)
        error.backward()
        opt.step()

Теперь, когда я бегу train_discriminator(discriminator,dataset,opt), я получаю Следующая ошибка, которую я не могу понять.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-15-ee20eb2a8e55> in <module>
----> 1 train_discriminator(discriminator,dataset,opt)

<ipython-input-13-9e6f9b4874c8> in train_discriminator(model, dataset, opt, n_iter, n_batch)
      3     for i in range(n_iter):
      4         X_real, y_real = generate_real_samples(dataset, half_batch)
----> 5         error = loss(model(X_real),y_real)
      6         error.backward()
      7         X_fake, y_fake = generate_fake_samples(half_batch)

~/environments/workspace/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    530             result = self._slow_forward(*input, **kwargs)
    531         else:
--> 532             result = self.forward(*input, **kwargs)
    533         for hook in self._forward_hooks.values():
    534             hook_result = hook(self, input, result)

~/environments/workspace/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    496 
    497     def forward(self, input, target):
--> 498         return F.binary_cross_entropy(input, target, weight=self.weight, reduction=self.reduction)
    499 
    500 

~/environments/workspace/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy(input, target, weight, size_average, reduce, reduction)
   2075 
   2076     return torch._C._nn.binary_cross_entropy(
-> 2077         input, target, weight, reduction_enum)
   2078 
   2079 

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' in call to _thnn_binary_cross_entropy_forward

Буду очень признателен, если кто-нибудь сможет предложить какие-либо изменения, которые необходимо внести, чтобы решить эту проблему.

Answer 1

Согласно сообщению об ошибке, основные истины не в графическом процессоре:

RuntimeError: Ожидаемый объект типа устройства cuda, но получил тип устройства cpu для аргумента # 2 'target' при вызове _thnn_binary_cross_entropy_forward