我收到 Float.Tensor 和 cuda.FloatTenson 不匹配的错误

I am getting this error of Float.Tensor and cuda.FloatTenson mismatch

我在 运行 模型的训练代码时遇到此错误。

Traceback (most recent call last):
  File "train.py", line 273, in <module>
    train_loss[epoch - 1] = process_epoch(
  File "train.py", line 240, in process_epoch
    loss = loss_fn(model, batch)
  File "train.py", line 221, in <lambda>
    loss_fn = lambda model, batch: weak_loss(model, batch, normalization="softmax")
  File "train.py", line 171, in weak_loss
    corr4d = model(batch).to("cuda")
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/srtf/ncnet/lib/model.py", line 263, in forward
    feature_A = self.FeatureExtraction(tnf_batch['source_image'])
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/srtf/ncnet/lib/model.py", line 84, in forward
    features = self.model(image_batch)
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/container.py", line 100, in forward
    input = module(input)
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 353, in forward
    return self._conv_forward(input, self.weight)
  File "/home/srtf/anaconda3/envs/ncnet/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 349, in _conv_forward
    return F.conv2d(input, weight, self.bias, self.stride,
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

系统上有 Cuda。我需要更改代码的哪些地方?

您的输入需要发送到正确的设备:

>>> corr4d = model(batch.cuda())

这会将批次复制到 GPU 设备(默认情况下 'cuda:0')。