Here’s my implementation just in case
class einsum_conv(nn.Module):
    def __init__(self, kernel_size):
        super(einsum_conv, self).__init__()
        self.ks = kernel_size

    def forward(self, x, kernel):
        if len(x.size()) == 3:
            x = x.unsqueeze(0)

        assert len(x.size()) == 4, 'need bs x c x h x w format'

        bs, in_c, h, w = x.size()
        ks = self.ks
        strided_x = x.as_strided((bs, in_c, h - ks + 1, w - ks + 1, ks, ks),
                                 (h * w * in_c, h * w, w, 1, w, 1))

        out = torch.einsum('bihwkl,oikl->bohw', (strided_x, kernel))
        return out