jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x):
        c2, c3, c4, c5 = x

        in5 = self.ins_conv[3](c5)
        in4 = self.ins_conv[2](c4)
        in3 = self.ins_conv[1](c3)
        in2 = self.ins_conv[0](c2)

        out4 = in4 + F.upsample(
            in5, scale_factor=2, mode="nearest")  # 1/16
        out3 = in3 + F.upsample(
            out4, scale_factor=2, mode="nearest")  # 1/8
        out2 = in2 + F.upsample(
            out3, scale_factor=2, mode="nearest")  # 1/4

        p5 = self.inp_conv[3](in5)
        p4 = self.inp_conv[2](out4)
        p3 = self.inp_conv[1](out3)
        p2 = self.inp_conv[0](out2)

        if self.intracl is True:
            p5 = self.incl4(p5)
            p4 = self.incl3(p4)
            p3 = self.incl2(p3)
            p2 = self.incl1(p2)

        p5 = F.upsample(p5, scale_factor=8, mode="nearest")
        p4 = F.upsample(p4, scale_factor=4, mode="nearest")
        p3 = F.upsample(p3, scale_factor=2, mode="nearest")

        fuse = torch.cat([p5, p4, p3, p2], dim=1)
        return fuse
//  [python]
def forward(self, x):
        c2, c3, c4, c5 = x

        in5 = self.ins_conv[3](c5)
        in4 = self.ins_conv[2](c4)
        in3 = self.ins_conv[1](c3)
        in2 = self.ins_conv[0](c2)

        out4 = in4 + F.interpolate(
            in5, scale_factor=2, mode="nearest")  # 1/16
        out3 = in3 + F.interpolate(
            out4, scale_factor=2, mode="nearest")  # 1/8
        out2 = in2 + F.interpolate(
            out3, scale_factor=2, mode="nearest")  # 1/4

        p5 = self.inp_conv[3](in5)
        p4 = self.inp_conv[2](out4)
        p3 = self.inp_conv[1](out3)
        p2 = self.inp_conv[0](out2)

        if self.intracl is True:
            p5 = self.incl4(p5)
            p4 = self.incl3(p4)
            p3 = self.incl2(p3)
            p2 = self.incl1(p2)

        p5 = F.interpolate(p5, scale_factor=8, mode="nearest")
        p4 = F.interpolate(p4, scale_factor=4, mode="nearest")
        p3 = F.interpolate(p3, scale_factor=2, mode="nearest")

        fuse = torch.cat([p5, p4, p3, p2], dim=1)
        return fuse
//  [python]
def forward(self, S, fout):
        W = self.permute(fout)
        size = int(2 ** self.level)
        M = (S > 1e-3).float()
        with torch.no_grad():
            maxW = F.max_pool2d((W + self.drift) * M, size, stride=[size, size]) - self.drift
            maxW = F.upsample_nearest(maxW, scale_factor=size) * M
        expW = torch.exp(W * M - maxW) * M
        avgS = F.avg_pool2d(S * expW, kernel_size=size, stride=size)
        avgexpW = F.avg_pool2d(expW, kernel_size=size, stride=size)
        Sp = avgS / (avgexpW + 1e-6)
        return Sp
//  [python]
def forward(self, S, fout):
        W = self.permute(fout)
        size = int(2 ** self.level)
        M = (S > 1e-3).float()
        with torch.no_grad():
            maxW = F.max_pool2d((W + self.drift) * M, size, stride=[size, size]) - self.drift
            maxW = F.interpolate(maxW, scale_factor=size, mode='nearest') * M
        expW = torch.exp(W * M - maxW) * M
        avgS = F.avg_pool2d(S * expW, kernel_size=size, stride=size)
        avgexpW = F.avg_pool2d(expW, kernel_size=size, stride=size)
        Sp = avgS / (avgexpW + 1e-6)
        return Sp
//  [python]
def forward(self, up_in):
        s = self.hook.stored
        up_out = self.shuf(up_in)
        ssh = s.shape[-2:]
        if ssh != up_out.shape[-2:]:
            up_out = F.upsample(up_out, size=s.shape[-2:], mode='nearest')
        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
        return self.conv2(self.conv1(cat_x))
//  [python]
def forward(self, up_in):
        s = self.hook.stored
        up_out = self.shuf(up_in)
        ssh = s.shape[-2:]
        if ssh != up_out.shape[-2:]:
            up_out = F.interpolate(up_out, s.shape[-2:], mode='nearest')
        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
        return self.conv2(self.conv1(cat_x))
//  [python]
def cropBox(img, ul, br, resH, resW):
    ul = ul.int()
    br = br.int()
    lenH = max(br[1] - ul[1], (br[0] - ul[0]) * resH / resW)
    lenW = lenH * resW / resH
    if img.dim() == 2:
        img = img[np.newaxis, :]

    newDim = torch.IntTensor((img.size(0), int(lenH), int(lenW)))
    newImg = img[:, ul[1]:, ul[0]:]
    # Crop and Padding
    size = torch.IntTensor((br[1] - ul[1], br[0] - ul[0]))

    newImg = SpecialCrop(size, 1)(newImg)
    newImg = Pad(newDim)(newImg)
    # Resize to output

    v_Img = torch.unsqueeze(newImg, 0)
    # newImg = F.upsample_bilinear(v_Img, size=(int(resH), int(resW))).data[0]
    newImg = F.upsample(v_Img, size=(int(resH), int(resW)),
                        mode='bilinear', align_corners=True).data[0]
    return newImg
//  [python]
def cropBox(img, ul, br, resH, resW):
    ul = ul.int()
    br = br.int()
    lenH = max(br[1] - ul[1], (br[0] - ul[0]) * resH / resW)
    lenW = lenH * resW / resH
    if img.dim() == 2:
        img = img[np.newaxis, :]

    newDim = torch.IntTensor((img.size(0), int(lenH), int(lenW)))
    newImg = img[:, ul[1]:, ul[0]:]
    # Crop and Padding
    size = torch.IntTensor((br[1] - ul[1], br[0] - ul[0]))

    newImg = SpecialCrop(size, 1)(newImg)
    newImg = Pad(newDim)(newImg)
    # Resize to output

    v_Img = torch.unsqueeze(newImg, 0)
    newImg = F.interpolate(v_Img, size=(int(resH), int(resW)),
                        mode='bilinear', align_corners=True).data[0]
    return newImg
//  [python]
def train(train_loader, net, criterion, optimizer, epoch, train_args):
    train_loss = AverageMeter()
    curr_iter = (epoch - 1) * len(train_loader)
    for i, data in enumerate(train_loader):
        #print(data)
        start = time.time()
        inputs, labels = data

        # test the input data
        # save_image = inputs[0]
        # save_label = labels[0]
        # plt.imsave('./i.png', (save_image.numpy().transpose(1, 2, 0) + 1) / 2)
        # plt.imsave('./l.png', (save_label.numpy() + 1) / 2)


        N = inputs.size(0)
        # to fit the origin operation
        inputs = Variable(inputs).cuda()
        labels = Variable(labels).cuda()
        #print(inputs)
        optimizer.zero_grad()
        outputs = net(inputs)

        # upsample 
        outputs = F.upsample_bilinear(outputs, size=(321,321))

        assert outputs.size()[2:] == labels.size()[1:]
        assert outputs.size()[1] == VOC.num_classes

        loss = criterion(outputs, labels) / N
        loss.backward()
        optimizer.step()

        train_loss.update(loss.data[0], N)

        curr_iter += 1

        if (i + 1) % train_args['print_freq'] == 0:
            end = time.time()
            i_time = end -start
            print('[epoch %d], [iter %d / %d], [train loss %.5f], [time %.5f per 10]' % (
                epoch, i + 1, len(train_loader), train_loss.avg, i_time*10
            ))
//  [python]
def train(train_loader, net, criterion, optimizer, epoch, train_args):
    train_loss = AverageMeter()
    curr_iter = (epoch - 1) * len(train_loader)
    for i, data in enumerate(train_loader):
        #print(data)
        start = time.time()
        inputs, labels = data

        # test the input data
        # save_image = inputs[0]
        # save_label = labels[0]
        # plt.imsave('./i.png', (save_image.numpy().transpose(1, 2, 0) + 1) / 2)
        # plt.imsave('./l.png', (save_label.numpy() + 1) / 2)


        N = inputs.size(0)
        # to fit the origin operation
        inputs = Variable(inputs).cuda()
        labels = Variable(labels).cuda()
        #print(inputs)
        optimizer.zero_grad()
        outputs = net(inputs)

        # upsample 
        outputs = F.interpolate(outputs, size=(321,321), mode='bilinear')

        assert outputs.size()[2:] == labels.size()[1:]
        assert outputs.size()[1] == VOC.num_classes

        loss = criterion(outputs, labels) / N
        loss.backward()
        optimizer.step()

        train_loss.update(loss.data[0], N)

        curr_iter += 1

        if (i + 1) % train_args['print_freq'] == 0:
            end = time.time()
            i_time = end -start
            print('[epoch %d], [iter %d / %d], [train loss %.5f], [time %.5f per 10]' % (
                epoch, i + 1, len(train_loader), train_loss.avg, i_time*10
            ))
//  [python]
def forward(self, x):
        # down sampling
        x = torch.nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
        x = torch.nn.functional.avg_pool2d(x, kernel_size=3, stride=1, padding=1)

        branch_0 = self.relu(self.norm_0(self.branch_0(x)))
        branch_1 = self.relu(self.norm_1(self.branch_1(x)))
        branch_2 = self.relu(self.norm_2(self.branch_2(x)))

        fusion = torch.cat([branch_0, branch_1, branch_2], 1)

        # score = torch.nn.functional.relu(self.softplus(self.attention_score(fusion)))
        score = self.softplus(self.attention_score(fusion))

        # recover resolution
        score = torch.nn.functional.upsample_nearest(score, scale_factor=2)
        
        return score
//  [python]
def forward(self, x):
        # down sampling
        x = torch.nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
        x = torch.nn.functional.avg_pool2d(x, kernel_size=3, stride=1, padding=1)

        branch_0 = self.relu(self.norm_0(self.branch_0(x)))
        branch_1 = self.relu(self.norm_1(self.branch_1(x)))
        branch_2 = self.relu(self.norm_2(self.branch_2(x)))

        fusion = torch.cat([branch_0, branch_1, branch_2], 1)

        # score = torch.nn.functional.relu(self.softplus(self.attention_score(fusion)))
        score = self.softplus(self.attention_score(fusion))

        # recover resolution
        score = torch.nn.functional.interpolate(score, scale_factor=2, mode='nearest')
        
        return score
//  [python]
def forward(self, x):
        out = F.upsample(x, scale_factor=(1, 2, 2), mode='nearest')
        out = self.conv(out)
        out = self.norm(out)
        out = F.relu(out)
        return out
//  [python]
def forward(self, x):
        out = F.interpolate(x, scale_factor=(1, 2, 2))
        out = self.conv(out)
        out = self.norm(out)
        out = F.relu(out)
        return out
//  [python]
def bilinear_interpolation(input: torch.Tensor, shape: List[int]) -> torch.Tensor:
    # input: B, H, W, C
    # shape: [target_height, target_width]
    return t.nn.functional.interpolate(input.permute(0, 3, 1, 2), size=shape, mode='bilinear', align_corners=False).permute(0, 2, 3, 1)
//  [python]
def bilinear_interpolation(input: torch.Tensor, shape: List[int]) -> torch.Tensor:
    # input: B, H, W, C
    # shape: [target_height, target_width]
    return F.interpolate(input.permute(0, 3, 1, 2), shape, mode='bilinear', align_corners=False).permute(0, 2, 3, 1)
//  [python]
def resize(clip, target_size, interpolation_mode):
    if len(target_size) != 2:
        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
    return torch.nn.functional.interpolate(clip, size=target_size, mode=interpolation_mode, align_corners=False)
//  [python]
def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x44 = self.up1(x5, x4)
        x33 = self.up2(x44, x3)
        x22 = self.up3(x33, x2)
        x11 = self.up4(x22, x1)
        x0 = self.outc(x11)
        x0 = self.sigmoid(x0)
        if self.deep_supervision:
            x11 = t.interpolate(
                self.dsoutc1(x11), size=x0.shape[2:], mode='bilinear')
            x22 = t.interpolate(
                self.dsoutc2(x22), size=x0.shape[2:], mode='bilinear')
            x33 = t.interpolate(
                self.dsoutc3(x33), size=x0.shape[2:], mode='bilinear')
            x44 = t.interpolate(
                self.dsoutc4(x44), size=x0.shape[2:], mode='bilinear')

            return x0, x11, x22, x33, x44
        else:
            return x0
//  [python]
def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x44 = self.up1(x5, x4)
        x33 = self.up2(x44, x3)
        x22 = self.up3(x33, x2)
        x11 = self.up4(x22, x1)
        x0 = self.outc(x11)
        x0 = self.sigmoid(x0)
        if self.deep_supervision:
            x11 = F.interpolate(
                self.dsoutc1(x11), x0.shape[2:], mode='bilinear')
            x22 = F.interpolate(
                self.dsoutc2(x22), x0.shape[2:], mode='bilinear')
            x33 = F.interpolate(
                self.dsoutc3(x33), x0.shape[2:], mode='bilinear')
            x44 = F.interpolate(
                self.dsoutc4(x44), x0.shape[2:], mode='bilinear')

            return x0, x11, x22, x33, x44
        else:
            return x0
//  [python]
def resize_scale(clip, target_size, interpolation_mode):
    if len(target_size) != 2:
        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
    H, W = clip.size(-2), clip.size(-1)
    scale_ = target_size[0] / min(H, W)
    return torch.nn.functional.interpolate(clip, scale_factor=scale_, mode=interpolation_mode, align_corners=False)
//  [python]
def forward(self, *xs):
        output = xs[0]

        if len(xs) == 2:
            output = self.skip_add.add(output, xs[1])

        output = t.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output
//  [python]
def forward(self, *xs):
        output = xs[0]

        if len(xs) == 2:
            output = self.skip_add.add(output, xs[1])

        output = nn.functional.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output
//  [python]
def get_abs_pos(abs_pos, has_cls_token, hw):
    h, w = hw
    if has_cls_token:
        abs_pos = abs_pos[:, 1:]
    xy_num = abs_pos.shape[1]
    size = int(math.sqrt(xy_num))
    assert size * size == xy_num

    if size != h or size != w:
        new_abs_pos = t.nn.functional.interpolate(
            abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2),
            size=(h, w),
            mode='bicubic',
            align_corners=False,
        )

        return new_abs_pos.permute(0, 2, 3, 1)
    else:
        return abs_pos.reshape(1, h, w, -1)
//  [python]
def get_abs_pos(abs_pos, has_cls_token, hw):
    h, w = hw
    if has_cls_token:
        abs_pos = abs_pos[:, 1:]
    xy_num = abs_pos.shape[1]
    size = int(math.sqrt(xy_num))
    assert size * size == xy_num

    if size != h or size != w:
        new_abs_pos = F.interpolate(
            abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2),
            size=(h, w),
            mode='bicubic',
            align_corners=False,
        )

        return new_abs_pos.permute(0, 2, 3, 1)
    else:
        return abs_pos.reshape(1, h, w, -1)
//  [python]
def get_rel_pos(q_size, k_size, rel_pos):
    """
    Get relative positional embeddings according to the relative positions
    of query and key sizes.
    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
        # Interpolate rel pos.
        rel_pos_resized = t.nn.functional.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode='linear',
        )
        rel_pos_resized = rel_pos_resized.reshape(-1,
                                                  max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords -
                       k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]
//  [python]
def get_rel_pos(q_size, k_size, rel_pos):
    """
    Get relative positional embeddings according to the relative positions
    of query and key sizes.
    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
        # Interpolate rel pos.
        rel_pos_resized = F.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode='linear',
        )
        rel_pos_resized = rel_pos_resized.reshape(-1,
                                                  max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords -
                       k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]
//  [python]
def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == self.num_ins
        outs = [inputs[0]]
        for i in range(1, self.num_ins):
            outs.append(
                t.interpolate(inputs[i], scale_factor=2**i, mode='bilinear'))
        out = torch.cat(outs, dim=1)
        if out.requires_grad and self.with_cp:
            out = checkpoint(self.reduction_conv, out)
        else:
            out = self.reduction_conv(out)
        outs = [out]
        for i in range(1, self.num_outs):
            outs.append(self.pooling(out, kernel_size=2**i, stride=2**i))
        outputs = []

        for i in range(self.num_outs):
            if outs[i].requires_grad and self.with_cp:
                tmp_out = checkpoint(self.fpn_convs[i], outs[i])
            else:
                tmp_out = self.fpn_convs[i](outs[i])
            outputs.append(tmp_out)
        return tuple(outputs)
//  [python]
def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == self.num_ins
        outs = [inputs[0]]
        for i in range(1, self.num_ins):
            outs.append(
                F.interpolate(inputs[i], scale_factor=2**i, mode='bilinear'))
        out = torch.cat(outs, dim=1)
        if out.requires_grad and self.with_cp:
            out = checkpoint(self.reduction_conv, out)
        else:
            out = self.reduction_conv(out)
        outs = [out]
        for i in range(1, self.num_outs):
            outs.append(self.pooling(out, kernel_size=2**i, stride=2**i))
        outputs = []

        for i in range(self.num_outs):
            if outs[i].requires_grad and self.with_cp:
                tmp_out = checkpoint(self.fpn_convs[i], outs[i])
            else:
                tmp_out = self.fpn_convs[i](outs[i])
            outputs.append(tmp_out)
        return tuple(outputs)
//  [python]
def forward(self, x):
        x = F.upsample(x, scale_factor=2, mode='bilinear', align_corners=False)
        return self.model(x)
//  [python]
def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
        return self.model(x)
//  [python]
def forward(self, hidden_states, output_size=None):
        assert hidden_states.shape[1] == self.channels

        if self.use_conv_transpose:
            raise NotImplementedError

        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
        dtype = hidden_states.dtype
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(torch.float32)

        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
        if hidden_states.shape[0] >= 64:
            hidden_states = hidden_states.contiguous()

        # if `output_size` is passed we force the interpolation output
        # size and do not make use of `scale_factor=2`
        if output_size is None:
            hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest")
        else:
            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")

        # If the input is bfloat16, we cast back to bfloat16
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype)

        hidden_states = self.conv(hidden_states)

        return hidden_states
//  [python]
def forward(self, feats: List[Tensor],
                batch_img_metas: List[dict]) -> Tuple[Tensor, Tensor]:
        """
        Args:
            feats (list[Tensor]): Feature maps of each level. Each has
                shape of (batch_size, c, h, w).
            batch_img_metas (list[dict]): List of image information.
                Pass in for creating more accurate padding mask. Not
                used here.

        Returns:
            tuple[Tensor, Tensor]: a tuple containing the following:

                - mask_feature (Tensor): Shape (batch_size, c, h, w).
                - memory (Tensor): Output of last stage of backbone.\
                        Shape (batch_size, c, h, w).
        """
        y = self.last_feat_conv(feats[-1])
        for i in range(self.num_inputs - 2, -1, -1):
            x = feats[i]
            cur_feat = self.lateral_convs[i](x)
            y = cur_feat + \
                t.nn.functional.interpolate(y, size=cur_feat.shape[-2:], mode='nearest')
            y = self.output_convs[i](y)

        mask_feature = self.mask_feature(y)
        memory = feats[-1]
        return mask_feature, memory
//  [python]
def forward(self, feats: List[Tensor],
                batch_img_metas: List[dict]) -> Tuple[Tensor, Tensor]:
        """
        Args:
            feats (list[Tensor]): Feature maps of each level. Each has
                shape of (batch_size, c, h, w).
            batch_img_metas (list[dict]): List of image information.
                Pass in for creating more accurate padding mask. Not
                used here.

        Returns:
            tuple[Tensor, Tensor]: a tuple containing the following:

                - mask_feature (Tensor): Shape (batch_size, c, h, w).
                - memory (Tensor): Output of last stage of backbone.\
                        Shape (batch_size, c, h, w).
        """
        y = self.last_feat_conv(feats[-1])
        for i in range(self.num_inputs - 2, -1, -1):
            x = feats[i]
            cur_feat = self.lateral_convs[i](x)
            y = cur_feat + \
                F.interpolate(y, size=cur_feat.shape[-2:], mode='nearest')
            y = self.output_convs[i](y)

        mask_feature = self.mask_feature(y)
        memory = feats[-1]
        return mask_feature, memory
//  [python]
def forward(self, xs):
        target_size = xs[-1].shape[2:]
        res = [F.upsample(x[self.idx[i]], size=target_size, mode='nearest') for i, x in enumerate(xs[:-1])]
        out = torch.sum(torch.stack(res + xs[-1:]), dim=0)
        return out
//  [python]
def forward(self, xs):
        target_size = xs[-1].shape[2:]
        res = [F.interpolate(x[self.idx[i]], size=target_size, mode='nearest') for i, x in enumerate(xs[:-1])]
        out = torch.sum(torch.stack(res + xs[-1:]), dim=0)
        return out
//  [python]
def __call__(self, x, skip=None):
        x = t.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
        if skip is not None:
            skip = spec_utils.crop_center(skip, x)
            x = torch.cat([x, skip], dim=1)
        h = self.conv(x)

        if self.dropout is not None:
            h = self.dropout(h)

        return h
//  [python]
def __call__(self, x, skip=None):
        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
        if skip is not None:
            skip = spec_utils.crop_center(skip, x)
            x = torch.cat([x, skip], dim=1)
        h = self.conv(x)

        if self.dropout is not None:
            h = self.dropout(h)

        return h
//  [python]
def imenlarge2(self, im):
        _, _, M, N = im.shape
        # t1 = F.upsample(im, size=(int(4*M-3), int(4*N-3)), mode='bilinear')
        t1 = F.interpolate(im, size=(int(4*M-3), int(4*N-3)), mode='bilinear', align_corners=False)
        t2 = torch.zeros([1, 1, 4*M-1, 4*N-1]).type(self.samplet.type())
        t2[:, :, 1: -1, 1:-1] = t1
        t2[:, :, 0, :] = 2*t2[:, :, 1, :] - t2[:, :, 2, :]
        t2[:, :, -1, :] = 2*t2[:, :, -2, :] - t2[:, :, -3, :]
        t2[:, :, :, 0] = 2*t2[:, :, :, 1] - t2[:, :, :, 2]
        t2[:, :, :, -1] = 2*t2[:, :, :, -2] - t2[:, :, :, -3]
        imu = t2[:, :, ::2, ::2]

        return imu
//  [python]
def _initialize_clip(config, num_patches=240):

        def resize_pos_embed(posemb, posemb_new):
            # Rescale the grid of position embeddings when loading from state_dict. Adapted from
            # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
            ntok_new = posemb_new.shape[1]

            posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
            ntok_new -= 1

            gs_old = int(math.sqrt(len(posemb_grid)))
            gs_new = int(math.sqrt(ntok_new))
            # _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
            posemb_grid = posemb_grid.reshape(1, gs_old, gs_old,
                                              -1).permute(0, 3, 1, 2)
            orig = posemb_grid.dtype
            posemb_grid = F.upsample(
                posemb_grid.float(), size=(gs_new, gs_new), mode='bilinear')
            posemb_grid = posemb_grid.to(orig)
            posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(
                1, gs_new * gs_new, -1)
            posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
            return posemb

        from .clip import clip
        clip_model = clip.load_from_config(config)
        if 'ViT-B-16' in config.clip_name:
            num_patches = int(config.image_res * config.image_res / (16 * 16))
            pos_embed = nn.Parameter(torch.zeros(num_patches + 1, 768).float())
        else:
            num_patches = int(config.image_res * config.image_res / (14 * 14))
            pos_embed = nn.Parameter(
                torch.zeros(num_patches + 1, 1024).float())
        pos_embed.weight = resize_pos_embed(
            clip_model.visual.positional_embedding.unsqueeze(0),
            pos_embed.unsqueeze(0))
        clip_model.visual.positional_embedding = pos_embed
        return clip_model
//  [python]
def _initialize_clip(config, num_patches=240):

        def resize_pos_embed(posemb, posemb_new):
            # Rescale the grid of position embeddings when loading from state_dict. Adapted from
            # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
            ntok_new = posemb_new.shape[1]

            posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
            ntok_new -= 1

            gs_old = int(math.sqrt(len(posemb_grid)))
            gs_new = int(math.sqrt(ntok_new))
            # _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
            posemb_grid = posemb_grid.reshape(1, gs_old, gs_old,
                                              -1).permute(0, 3, 1, 2)
            orig = posemb_grid.dtype
            posemb_grid = F.interpolate(
                posemb_grid.float(), size=(gs_new, gs_new), mode='bilinear')
            posemb_grid = posemb_grid.to(orig)
            posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(
                1, gs_new * gs_new, -1)
            posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
            return posemb

        from .clip import clip
        clip_model = clip.load_from_config(config)
        if 'ViT-B-16' in config.clip_name:
            num_patches = int(config.image_res * config.image_res / (16 * 16))
            pos_embed = nn.Parameter(torch.zeros(num_patches + 1, 768).float())
        else:
            num_patches = int(config.image_res * config.image_res / (14 * 14))
            pos_embed = nn.Parameter(
                torch.zeros(num_patches + 1, 1024).float())
        pos_embed.weight = resize_pos_embed(
            clip_model.visual.positional_embedding.unsqueeze(0),
            pos_embed.unsqueeze(0))
        clip_model.visual.positional_embedding = pos_embed
        return clip_model
//  [python]
def forward(self, x):
        _, _, h, w = x.size()
        feat1 = F.upsample(
            self.conv1(x), size=(h, w), mode="bilinear"
        )
        feat2 = self.conv2(x)
        feat3 = self.conv3(x)
        feat4 = self.conv4(x)
        feat5 = self.conv5(x)
        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
        bottle = self.bottleneck(out)
        return bottle
//  [python]
def forward(self, x):
        _, _, h, w = x.size()
        feat1 = F.interpolate(
            self.conv1(x), size=(h, w), mode="bilinear", align_corners=True
        )
        feat2 = self.conv2(x)
        feat3 = self.conv3(x)
        feat4 = self.conv4(x)
        feat5 = self.conv5(x)
        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
        bottle = self.bottleneck(out)
        return bottle
//  [python]
def forward(self, x):
        out = x
        for i in range(0, len(self.body)):
            out = self.body[i](out)

        out = self.upsampler(out)
        # add the nearest upsampled image, so that the network learns the residual
        base = F.upsample(x, scale_factor=self.scale, mode="nearest")
        out += base
        return out
//  [python]
def forward(self, x):
        out = x
        for i in range(0, len(self.body)):
            out = self.body[i](out)

        out = self.upsampler(out)
        # add the nearest upsampled image, so that the network learns the residual
        base = F.interpolate(x, scale_factor=self.scale, mode="nearest")
        out += base
        return out
//  [python]
def _fuse_feature(x, feature):
        _, _, h, w = feature.shape
        x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)
        x = torch.cat([x, feature], dim=1)
        return x
//  [python]
def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, C, F, T).

        Returns:
            Tensor: Interpolated tensor (B, C, F * y_scale, T * x_scale),

        """
        return F.interpolate(x, size=(int(x.shape[2] * self.y_scale), int(x.shape[3] * self.x_scale)), mode=self.mode)
//  [python]
def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, C, F, T).

        Returns:
            Tensor: Interpolated tensor (B, C, F * y_scale, T * x_scale),

        """
        return F.interpolate(x, scale_factor=(self.y_scale, self.x_scale), mode=self.mode)
//  [python]
def preprocess_single(video, resolution=224, sequence_length=None):
    # video: CTHW, [0, 1]
    c, t, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:, :sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = t.interpolate(video, size=target_size, mode='bilinear', align_corners=False)

    # center crop
    c, t, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]

    # [0, 1] -> [-1, 1]
    video = (video - 0.5) * 2
    return video.contiguous()
//  [python]
def preprocess_single(video, resolution=224, sequence_length=None):
    # video: CTHW, [0, 1]
    c, t, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:, :sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False)

    # center crop
    c, t, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]

    # [0, 1] -> [-1, 1]
    video = (video - 0.5) * 2
    return video.contiguous()
//  [python]
def forward(self, x, up_x):
        out = self.in_skip(x)

        if self.ind > 0:
            up_x = F.upsample(
                up_x,
                size=(x.size(2), x.size(3)),
                mode='bilinear')
            up_x = self.up_conv(up_x)
            out = out + up_x
        out = self.relu(out)

        skip1 = None
        skip2 = None
        if self.gen_skip:
            skip1 = self.out_skip1(x)
            skip2 = self.out_skip2(out)

        cross_conv = None
        if self.ind == self.num_units - 1 and self.gen_cross_conv:
            cross_conv = self.cross_conv(out)

        return out, skip1, skip2, cross_conv
//  [python]
def forward(self, x, up_x):
        out = self.in_skip(x)

        if self.ind > 0:
            up_x = F.interpolate(
                up_x,
                size=(x.size(2), x.size(3)),
                mode='bilinear',
                align_corners=True)
            up_x = self.up_conv(up_x)
            out = out + up_x
        out = self.relu(out)

        skip1 = None
        skip2 = None
        if self.gen_skip:
            skip1 = self.out_skip1(x)
            skip2 = self.out_skip2(out)

        cross_conv = None
        if self.ind == self.num_units - 1 and self.gen_cross_conv:
            cross_conv = self.cross_conv(out)

        return out, skip1, skip2, cross_conv
//  [python]
def forward(self, x):
        result_dict = {}
        layer_names = ['conv1_1',
                       'conv2_1',
                       'conv3_1', 'conv3_2',
                       'conv4_1', 'conv4_2',
                       'conv5_1', 'conv5_2']
        idx = 0
        image_in = x.clone()
        for ii, model in enumerate(self.features):
            x = model(x)
            if ii in {0, 3, 6, 8, 11, 13, 16, 18}:
                result_dict[layer_names[idx]] = x
                idx += 1

        # pdb.set_trace()
        # merge different level features of VGG
        merge_feature_names = ['conv1_1', 'conv3_2', 'conv5_2']
        out_feature = None
        _, c, h, w = image_in.shape
        for mfn in merge_feature_names:
            # resize all the extracted feature maps to img size, and concat them all
            feat = result_dict[mfn]
            feat = F.upsample(feat, size=(h, w), mode='bilinear')
            out_feature = torch.cat([out_feature, feat], dim=1) if out_feature is not None else feat

        # pdb.set_trace()
        out_feature = self.merge_features(out_feature)
        return out_feature
//  [python]
def forward(self, x):
        result_dict = {}
        layer_names = ['conv1_1',
                       'conv2_1',
                       'conv3_1', 'conv3_2',
                       'conv4_1', 'conv4_2',
                       'conv5_1', 'conv5_2']
        idx = 0
        image_in = x.clone()
        for ii, model in enumerate(self.features):
            x = model(x)
            if ii in {0, 3, 6, 8, 11, 13, 16, 18}:
                result_dict[layer_names[idx]] = x
                idx += 1

        # pdb.set_trace()
        # merge different level features of VGG
        merge_feature_names = ['conv1_1', 'conv3_2', 'conv5_2']
        out_feature = None
        _, c, h, w = image_in.shape
        for mfn in merge_feature_names:
            # resize all the extracted feature maps to img size, and concat them all
            feat = result_dict[mfn]
            feat = F.interpolate(feat, size=(h, w), mode='bilinear', align_corners=True)
            out_feature = torch.cat([out_feature, feat], dim=1) if out_feature is not None else feat

        # pdb.set_trace()
        out_feature = self.merge_features(out_feature)
        return out_feature
//  [python]
def forward(self, masked_frames, masks):
        b, t, c, h, w = masked_frames.size()
        masks = masks.view(b * t, 1, h, w)
        enc_feat = self.encoder(masked_frames.view(b * t, c, h, w))
        _, c, h, w = enc_feat.size()
        masks = F.upsample(masks, scale_factor=1.0 / 4, mode='nearest')
        enc_feat = self.transformer({
            'x': enc_feat,
            'm': masks,
            'b': b,
            'c': c
        })['x']
        output = self.decoder(enc_feat)
        output = torch.tanh(output)
        return output
//  [python]
def forward(self, masked_frames, masks):
        b, t, c, h, w = masked_frames.size()
        masks = masks.view(b * t, 1, h, w)
        enc_feat = self.encoder(masked_frames.view(b * t, c, h, w))
        _, c, h, w = enc_feat.size()
        masks = F.interpolate(masks, scale_factor=1.0 / 4)
        enc_feat = self.transformer({
            'x': enc_feat,
            'm': masks,
            'b': b,
            'c': c
        })['x']
        output = self.decoder(enc_feat)
        output = torch.tanh(output)
        return output
//  [python]
def forward(self, x):
        result_dict = {}
        layer_names = ['conv1_1', 'conv1_2',
                       'conv2_1', 'conv2_2',
                       'conv3_1', 'conv3_2', 'conv3_3', 'conv3_4',
                       'conv4_1', 'conv4_2', 'conv4_3', 'conv4_4',
                       'conv5_1', 'conv5_2', 'conv5_3', 'conv5_4']
        idx = 0
        image_in = x.clone()
        for ii, model in enumerate(self.features):
            x = model(x)
            if ii in {0, 2, 5, 7, 10, 12, 14, 16, 19, 21, 23, 25, 28, 30, 32, 34}:
                result_dict[layer_names[idx]] = x
                idx += 1

        # merge different level features of VGG
        merge_feature_names = ['conv1_2', 'conv3_4', 'conv5_4']
        out_feature = None
        _, c, h, w = image_in.shape
        for mfn in merge_feature_names:
            feat = result_dict[mfn]
            feat = F.upsample(feat, size=(h, w), mode='bilinear')
            out_feature = torch.cat([out_feature, feat], dim=1) if out_feature is not None else feat

        out_feature = self.merge_features(out_feature)
        return out_feature
//  [python]
def forward(self, x):
        result_dict = {}
        layer_names = ['conv1_1', 'conv1_2',
                       'conv2_1', 'conv2_2',
                       'conv3_1', 'conv3_2', 'conv3_3', 'conv3_4',
                       'conv4_1', 'conv4_2', 'conv4_3', 'conv4_4',
                       'conv5_1', 'conv5_2', 'conv5_3', 'conv5_4']
        idx = 0
        image_in = x.clone()
        for ii, model in enumerate(self.features):
            x = model(x)
            if ii in {0, 2, 5, 7, 10, 12, 14, 16, 19, 21, 23, 25, 28, 30, 32, 34}:
                result_dict[layer_names[idx]] = x
                idx += 1

        # merge different level features of VGG
        merge_feature_names = ['conv1_2', 'conv3_4', 'conv5_4']
        out_feature = None
        _, c, h, w = image_in.shape
        for mfn in merge_feature_names:
            feat = result_dict[mfn]
            feat = F.interpolate(feat, size=(h, w), mode='bilinear', align_corners=True)
            out_feature = torch.cat([out_feature, feat], dim=1) if out_feature is not None else feat

        out_feature = self.merge_features(out_feature)
        return out_feature
//  [python]
def infer(self, feat, masks):
        t, c, h, w = masks.size()
        masks = masks.view(t, c, h, w)
        masks = F.upsample(masks, scale_factor=1.0 / 4, mode='nearest')
        t, c, _, _ = feat.size()
        enc_feat = self.transformer({
            'x': feat,
            'm': masks,
            'b': 1,
            'c': c
        })['x']
        return enc_feat
//  [python]
def infer(self, feat, masks):
        t, c, h, w = masks.size()
        masks = masks.view(t, c, h, w)
        masks = F.interpolate(masks, scale_factor=1.0 / 4)
        t, c, _, _ = feat.size()
        enc_feat = self.transformer({
            'x': feat,
            'm': masks,
            'b': 1,
            'c': c
        })['x']
        return enc_feat
//  [python]
def forward(self, x):
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.lrelu(self.conv_up2(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x
//  [python]
def forward(self, x):
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x
//  [python]
def upsample(x, h, w):
    return t.interpolate(x, size=[h, w], mode='bicubic', align_corners=True)
//  [python]
def upsample(x, h, w):
    return F.interpolate(x, size=[h, w], mode='bicubic', align_corners=True)
//  [python]
def forward(self, pred_flows, gt_local_frames):
        b, l_t, c, h, w = gt_local_frames.size()

        with torch.no_grad():
            # compute gt forward and backward flows
            gt_local_frames = t.interpolate(gt_local_frames.view(-1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True)
            gt_local_frames = gt_local_frames.view(b, l_t, c, h // 4, w // 4)
            gtlf_1 = gt_local_frames[:, :-1, :, :, :].reshape(
                -1, c, h // 4, w // 4)
            gtlf_2 = gt_local_frames[:, 1:, :, :, :].reshape(
                -1, c, h // 4, w // 4)
            gt_flows_forward = self.fix_spynet(gtlf_1, gtlf_2)
            gt_flows_backward = self.fix_spynet(gtlf_2, gtlf_1)

        # calculate loss for flow completion
        forward_flow_loss = self.l1_criterion(
            pred_flows[0].view(-1, 2, h // 4, w // 4), gt_flows_forward)
        backward_flow_loss = self.l1_criterion(
            pred_flows[1].view(-1, 2, h // 4, w // 4), gt_flows_backward)
        flow_loss = forward_flow_loss + backward_flow_loss

        return flow_loss
//  [python]
def forward(self, pred_flows, gt_local_frames):
        b, l_t, c, h, w = gt_local_frames.size()

        with torch.no_grad():
            # compute gt forward and backward flows
            gt_local_frames = F.interpolate(gt_local_frames.view(-1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True,
                                            recompute_scale_factor=True)
            gt_local_frames = gt_local_frames.view(b, l_t, c, h // 4, w // 4)
            gtlf_1 = gt_local_frames[:, :-1, :, :, :].reshape(
                -1, c, h // 4, w // 4)
            gtlf_2 = gt_local_frames[:, 1:, :, :, :].reshape(
                -1, c, h // 4, w // 4)
            gt_flows_forward = self.fix_spynet(gtlf_1, gtlf_2)
            gt_flows_backward = self.fix_spynet(gtlf_2, gtlf_1)

        # calculate loss for flow completion
        forward_flow_loss = self.l1_criterion(
            pred_flows[0].view(-1, 2, h // 4, w // 4), gt_flows_forward)
        backward_flow_loss = self.l1_criterion(
            pred_flows[1].view(-1, 2, h // 4, w // 4), gt_flows_backward)
        flow_loss = forward_flow_loss + backward_flow_loss

        return flow_loss
//  [python]
def cat(self, tensors, *args, **kwargs):
        if len(tensors) == 2:
            a, b = tensors
            if a.shape[-2:] != b.shape[-2:]:
                a = t.nn.functional.interpolate(a, b.shape[-2:], mode="nearest")

            tensors = (a, b)

        return torch.cat(tensors, *args, **kwargs)
//  [python]
def cat(self, tensors, *args, **kwargs):
        if len(tensors) == 2:
            a, b = tensors
            if a.shape[-2:] != b.shape[-2:]:
                a = torch.nn.functional.interpolate(a, b.shape[-2:], mode="nearest")

            tensors = (a, b)

        return torch.cat(tensors, *args, **kwargs)
//  [python]
def upsample_pytorch_test(self, h, w, scale_h, scale_w, align_corners, cpu_only):
        input_dim = (1, 1, h, w)
        if align_corners:
            linear_upsample_mode = "ALIGN_CORNERS_TRUE"
        else:
            linear_upsample_mode = "ALIGN_CORNERS_FALSE"

        input_features = [("data", datatypes.Array(*input_dim))]
        output_features = [("output", None)]

        builder = neural_network.NeuralNetworkBuilder(
            input_features, output_features, disable_rank5_shape_mapping=True
        )
        builder.add_upsample(
            name="upsample",
            scaling_factor_h=scale_h,
            scaling_factor_w=scale_w,
            linear_upsample_mode=linear_upsample_mode,
            input_name="data",
            output_name="output",
            mode="BILINEAR",
        )

        input_tensor = np.reshape(np.arange(1.0, 1.0 + (h * w), 1.0), input_dim)
        input = {"data": input_tensor}

        # Get result from PyTorch
        x = torch.from_numpy(np.reshape(input_tensor, (1, 1, h, w)))
        pytorch_output = torch.nn.functional.upsample_bilinear(
            x,
            size=(int(h * scale_h), int(w * scale_w)),
        )

        # Expect PyTorch output matches CoreML output
        expected = {"output": pytorch_output.numpy()}

        self._test_model(builder.spec, input, expected, useCPUOnly=cpu_only)
        self.assertEqual(len(input_dim), builder._get_rank("output"))