This code should work. For some reasons I don’t know, the community don’t like to add this function into the core implementation.
def squeeze(x, r):
    [B, C, H, W] = list(x.size())
    x = x.reshape(B, C, H//r, r, W//r, r)
    x = x.permute(0, 1, 3, 5, 2, 4)
    x = x.reshape(B, C*(r**2), H//r, W//r)
    return x