I think the main difficulty stems from the use of nn.Sequential, which only takes in a single argument input. So it’s not possible to do branching in nn.Sequential when additional info is required (in this case, the data type). Which means that we need to replace all nn.Sequential.
For BasicBlock:
class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)

        # self.shortcut = nn.Sequential()
        # if stride != 1 or in_planes != self.expansion*planes:
        #     self.shortcut = nn.Sequential(
        #         nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
        #         nn.BatchNorm2d(self.expansion*planes)
        #     )
        if stride != 1 or in_planes != self.expansion*planes:
            self.identity_shortcut = False
            self.shortcut_conv = nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)
            self.shortcut_bn1 = nn.BatchNorm2d(self.expansion*planes)
            self.shortcut_bn2 = nn.BatchNorm2d(self.expansion*planes)
        else:
            self.identity_shortcut = True

    def forward(self, x, data_type):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))

        # out += self.shortcut(x)
        if not self.identity_shortcut:
            x = self.shortcut_conv(x)
            if data_type == 'type1':
                x = self.sortcut_bn1(x)
            elif data_type == 'type2':
                x = self.sortcut_bn2(x)
        out += x

        out = F.relu(out)
        return out

For ResNet._make_layer:
def _make_layer(self, block, planes, num_blocks, stride):
    strides = [stride] + [1]*(num_blocks-1)
    layers = []
    for stride in strides:
        layers.append(block(self.in_planes, planes, stride))
        self.in_planes = planes * block.expansion
    return nn.ModuleList(layers)

and ResNet.forward:
def forward(self, x, data_type):
    out = F.relu(self.bn1(self.conv1(x)))
    # out = self.layer1(out)
    # out = self.layer2(out)
    # out = self.layer3(out)
    # out = self.layer4(out)
    for layer in self.layer1:
        out = layer(out, data_type)
    for layer in self.layer2:
        out = layer(out, data_type)
    for layer in self.layer3:
        out = layer(out, data_type)
    for layer in self.layer4:
        out = layer(out, data_type)
    out = F.avg_pool2d(out, 4)
    out = out.view(out.size(0), -1)
    out = self.linear(out)
    return out

This is only a quick fix, I’m sure the code can be made more elegant.