7c3d0a5ea7405fe74cc31f5553f7c04d9804d42e,pytorch/pytorchcv/models/isqrtcovresnet.py,Covpool,forward,#Any#Any#,21

Before Change

    def forward(ctx, input):
        x = input
        batch_size = x.data.shape[0]
        channels = x.data.shape[1]
        h = x.data.shape[2]
        w = x.data.shape[3]
        M = h * w
        x = x.reshape(batch_size, channels, M)

After Change

        batch, channels, height, width = x.size()
        n = height * width
        xn = x.reshape(batch, channels, n)
        identity_bar = ((1.0 / n) * torch.eye(n, dtype=xn.dtype, device=xn.device)).unsqueeze(dim=0).repeat(batch, 1, 1)
        ones_bar = torch.full((batch, n, n), fill_value=(-1.0 / n / n), dtype=xn.dtype, device=xn.device)
        i_bar = identity_bar + ones_bar
        sigma = xn.bmm(i_bar).bmm(xn.transpose(1, 2))
        ctx.save_for_backward(x, i_bar)
        return sigma

Instances