You can implement it as an extension of torch.nn.utils.prune. You can find instructions on how to create your custom pruning function by following the pytorch pruning tutorial: https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#extending-torch-nn-utils-prune-with-custom-pruning-functions 60
In practice, in compute_mask, among other things, you will have to generate a mask that preserves all the entries in the tensor t that are above your threshold: mask = (t > 1.).to(t.dtype)