Efficient way of doing this is having byte_tensor mask with same size as the original .weight tensor. So you can just remove .float() casting. Your mask would have ones for pruned weights. You don’t need a Variable of masks. You can just use the masks for indexing like below for example.
loss.backward()
for layer,mask in zip(X,Y):
     layer.weight.grad.data[mask]=0
optimizer.step()

You just need to keep track which mask belongs to which weight. That’s MaskedModule wrapper is doing in the code I shared above.