How to make custom method in nn.Module work with GPUs
I’m trying to implement simple res-net like below and it works with CPU. <SCODE>class ResNet(nn.module):
      def __init__(self):
        super(Net, self).__init__()
        self.conv_1 = nn.Conv2d(3, 64, 4, stride=2)
        self.bn_1 = nn.BatchNorm2d(64)
        self.res_1 = self.__res_block(64, [32,32,128], True)
         ...
    def forward(self, x):
        x = self.conv_1(x)
        x = F.relu(self.bn_1(x))
        x = F.max_pool2d(x, 2, 2)
        x = self.res_1(x)
        ...

    def __res_block(self, in_channels,
                    nb_filters, right=False):
            def __res_base(_in_channels, out_channels,
                       kernel_size=1, padding=0):
            def g(x):
                x = nn.Conv2d(in_channels=_in_channels,
                              out_channels=out_channels,
                              kernel_size=kernel_size,
                              padding=padding)(x)
                x = nn.BatchNorm2d(num_features=out_channels)(x)
                return x
            return g

        def f(x):
            y = F.relu(__res_base(in_channels, nb_filters[0])(x))
            y = F.relu(__res_base(nb_filters[0], nb_filters[1],
                                  kernel_size=3, padding=1)(y))
            y = F.relu(__res_base(nb_filters[1], nb_filters[2])(y))
            if right is True:
                x = __res_base(in_channels, nb_filters[2])(x)
            return F.relu(x+y)
        return f
<ECODE> but it doesn’t work with GPU and throws TypeError, <SCODE>TypeError: _cudnn_convolution_full_forward received an invalid combination of arguments 
- got (torch.cuda.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.cuda.FloatTensor, tuple, tuple, int, bool), 
but expected (torch.cuda.RealTensor input, torch.cuda.RealTensor weight, torch.cuda.RealTensor bias, 
torch.cuda.RealTensor output, std::vector<int> pad, std::vector<int> stride, int groups, bool benchmark)
<ECODE>
moskomule
Hi,
albanD
moskomule
alexis-jacq
Ho I see,
albanD
apaszke
moskomule
smth
moskomule
apaszke
marioviti