I think this is the snippet that is giving the error. If i dont use this snippet then the model works in multi-gpu
class RcnnMultiHead(nn.Module):
    def __init__(self, scales, cfg, in_channels):
        super(RcnnMultiHead, self).__init__()
        self.scales = scales
        self.num_classes = cfg.num_classes
        self.crop_size = cfg.rcnn_crop_size
        
        self.fc1s = nn.ModuleList()
        self.fc2s = nn.ModuleList()
        self.logits = nn.ModuleList()
        self.deltas = nn.ModuleList()
        for i in range(self.scales):
            self.fc1s.append(nn.Linear(in_channels * cfg.rcnn_scales[i] * cfg.rcnn_scales[i], 784))
            self.fc2s.append(nn.Linear(784, 784))
            self.logits.append(nn.Linear(784, self.num_classes))
            self.deltas.append(nn.Linear(784, self.num_classes * 4))
    
    def forward(self, crops):
        logits_flat = []
        deltas_flat = []
        for i in range(self.scales):
            crop = crops[i]
            x = crop.view(crop.size(0), -1)
            x = F.relu(self.fc1s[i](x), inplace=True)
            x = F.relu(self.fc2s[i](x), inplace=True)
            x = F.dropout(x, 0.5, training=self.training)
            logit = self.logits[i](x)
            delta = self.deltas[i](x)
            logits_flat.append(logit)
            deltas_flat.append(delta)
        logits_flat = torch.cat(logits_flat, 0)
        deltas_flat = torch.cat(deltas_flat, 0)
        return logits_flat, deltas_flat

class RcnnFpn(nn.Module):
    def __init__(self, in_dim):
        super(RcnnFpn, self).__init__()
        self.in_dim = in_dim
        self.c1 = nn.Conv2d(256, 384, kernel_size=3, stride=2, padding=1)
        self.p1 = nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0)
        self.c1_2 = LateralBlock(384, 256, self.in_dim)
        self.c2 = nn.Conv2d(384, 512, kernel_size=3, stride=2, padding=1)
        self.c2_2 = LateralBlock(256, self.in_dim, self.in_dim)
    
    def forward(self, x):
        c1_out = F.leaky_relu(self.c1(x))
        c2_out = F.leaky_relu(self.c2(c1_out))
        p1 = self.p1(c2_out)
        p2 = self.c1_2(c1_out, p1)
        p3 = self.c2_2(x, p2)
        features = [p1, p2, p3]
        return features

class Model(nn.Module):
    def __init__(self, cfg):
        super(Model, self).__init__()
        .
        .
        .
        self.rcnn_crop = CropRoi(cfg, cfg.rcnn_crop_size)
        self.rcnn_head = RcnnMultiHead(3, cfg, crop_channels)
        self.rcnn_fpn = RcnnFpn(in_dim=256)
    def forward(self, x):
        .
        .
        .
        rcnn_crops = self.rcnn_crop(features, self.rpn_proposals)
        rcnn_features = data_parallel(self.rcnn_fpn, rcnn_crops)
        self.rcnn_logits, self.rcnn_deltas = data_parallel(self.rcnn_head, rcnn_features)  

The self.rcnn_fpn and self.rcnn_head is called by data_parallel function from nn.parallel.data_parallel and it is causing the error.