pruning

[starsky68]

class ChooseTopEdges(autograd.Function):
""" Chooses the top edges for the forwards pass but allows gradient flow to all edges in the backwards pass"""

@staticmethod
def forward(ctx, weight, prune_rate):
    output = weight.clone()
    _, idx = weight.flatten().abs().sort()
    p = int(prune_rate * weight.numel())
    flat_oup = output.flatten()
    flat_oup[idx[:p]] = 0
    return output

@staticmethod
def backward(ctx, grad_output):
    return grad_output, None, None

Excuse me, although flat_oup is pruned because the element is set to zero, the return in forward is output. Does this play the role of pruning? @dirkgr @schmmd @iellenberger @danyaljj @danyaljj @MLatzke