I am trying to implement Group Lasso on weight matrices of a neural network in PyTorch.
I have written the code to implement Group Lasso but am unsure if this is correct, confirmation or correction of my code will be very helpful.
def gl_norm(model, gl_lambda, num_blk):
    gl_reg = torch.tensor(0., dtype=torch.float32).cuda()
    for key in model:
        for param in model[key].parameters():
            dim = param.size()
            if dim.__len__() > 1 and not model[key].skip_regularization:
                div1 = list(torch.chunk(param,int(num_blk),1))
                all_blks = []
                for div2 in div1:
                    temp = list(torch.chunk(div2,int(num_blk),0))
                    for blk in temp:
                        all_blks.append(blk)
                for l2_param in all_blks:
                    gl_reg += torch.norm(l2_param, 2)
    return gl_reg * float(gl_lambda)

I expect the torch.chunk function to break up the weight matrix into small blocks which then go through L2 norm for the block and L1 norm between all the blocks.