Basic implementation of maxpool2d in pytorch, based on it you can create kmaxpool

Pytorch_maxpool2d.py

stride = 3
new_var = Variable(torch.zeros([x.shape[0], x.shape[1]//stride, x.shape[2]//stride]))
for dim1 in range(x.shape[0]):
    tmp = Variable(torch.zeros([x.shape[1]//stride, x.shape[2]//stride, 1]))
    for i in range(0, x.shape[1], stride):
        for j in range(0, x.shape[2], stride):
            tmp_max = x[dim1][i][j]
            for k in range(stride):
                for m in range(stride):
                    tmp_max = torch.max(tmp_max, x[dim1][i+k][j+m])
            tmp[i//stride, j//stride] = tmp_max
    new_var[dim1] = tmp