pytorch 踩坑

代码如下

........

 self.features = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )

.......

x = x / 255.

x = self.features(x)

遇到错误：

Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument.........

因为第一个x是在cpu上的，而features(x)返回一个gpu上的变量，需要加上类型转化：

x = x / 255.

x = x.type(torch.cuda.FloatTensor)

x = self.features(x)

这样，即可通过。