//  [python]
def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.layer1(x)
        x_head_1 = x

        x_list = []
        for i in range(self.stage2_cfg['NUM_BRANCHES']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)
        y_list = self.stage2(x_list)

        x_list = []
        for i in range(self.stage3_cfg['NUM_BRANCHES']):
            if self.transition2[i] is not None:
                x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])

        x = self.stage3(x_list)

        # Replace the classification heaeder with custom setting
        # Upsampling
        x0_h, x0_w = x[0].size(2), x[0].size(3)
        x1 = F.interpolate(
            x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=False)
        x2 = F.interpolate(
            x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=False)
        x = torch.cat([x[0], x1, x2, x_head_1], 1)
        # first, reduce the channel down
        x = self.redc_layer(x)

        pred_attn = self.GAP(F.relu_(self.pred_bn(self.pred_conv(x))))
        pred_attn = F.softmax(pred_attn, dim=1)
        pred_attn_list = torch.chunk(pred_attn, 4, dim=1)

        aspp_out = []
        for k, v in enumerate(self.aspp):
            if k % 2 == 0:
                aspp_out.append(self.aspp[k + 1](v(x)))
            else:
                continue
        # Using Aspp add, and relu inside
        for i in range(4):
            x = x + F.relu_(aspp_out[i] * 0.25) * pred_attn_list[i]

        bz = x.size(0)
        # -- Besides, we also need to let the prediction attention be close to visable domain
        # -- Calculate the domain distance and get the weights
        # - First, detach domains
        G_all_d = self.G_all.detach()  # use detached G_all for calulcating
        pred_attn_d = pred_attn.detach().view(bz, 512, 1, 1)

        if self.cosine == 1:
            G_A, G_B, G_Q = torch.chunk(G_all_d, self.n_domain, dim=0)

            cos_dis_A = F.cosine_similarity(pred_attn_d, G_A, dim=1).view(-1)
            cos_dis_B = F.cosine_similarity(pred_attn_d, G_B, dim=1).view(-1)
            cos_dis_Q = F.cosine_similarity(pred_attn_d, G_Q, dim=1).view(-1)

            cos_dis_all = torch.stack([cos_dis_A, cos_dis_B,
                                       cos_dis_Q]).view(bz, -1)  # bz*3

            cos_dis_all = F.softmax(cos_dis_all, dim=1)

            target_attn = cos_dis_all.view(bz, self.n_domain, 1, 1, 1).expand(
                bz, self.n_domain, 512, 1, 1) * self.G_all.view(
                    1, self.n_domain, 512, 1, 1).expand(
                        bz, self.n_domain, 512, 1, 1)
            target_attn = torch.sum(
                target_attn, dim=1, keepdim=False)  # bz * 512 * 1 * 1

            if self.fix_domain:
                target_attn = target_attn.detach()

        else:
            raise ValueError('Have not implemented not cosine distance yet')

        x = self.last_layer(x)
        x = F.relu_(x)

        x = F.interpolate(
            x, size=(x0_h * 2, x0_w * 2), mode='bilinear', align_corners=False)

        return x, pred_attn, target_attn