Hi @spanev
I really appreciate for your help, but I got another one  
def hook(module, input):
    with torch.no_grad():
        module.weight.data = module.weight.data * self.sparse_mask[self.type[module]]

I used exactly what you show me, and I got this error

RuntimeError: expected device cuda:0 and dtype Float but got device cpu and dtype Float

Could you please teach me how to fix this problem?
Any response will be appreciated!