RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found

RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu

nn.Sequential(
                nn.Conv2d(3, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(512), nn.ReLU(inplace=True),
                nn.AdaptiveAvgPool2d((1, 1)),
                nn.Flatten(),
                )

在模型最后加.cuda()
改为:

nn.Sequential(
                nn.Conv2d(3, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(512), nn.ReLU(inplace=True),
                nn.AdaptiveAvgPool2d((1, 1)),
                nn.Flatten(),
                ).cuda()

你可能感兴趣的:(Python,深度学习,python,人工智能)