在训练Pytorch模型时，如何掩盖部分参数？

我正在训练一个像往常一样初始化的简单pytorch模型：

           model = nn.Sequential(
            nn.Linear(input_dim,output_dim)).

我还有一个特定的矩阵W，它的大小为（input_dim，output_dim），该矩阵W由零和一个我想用作整个训练和评分过程中的掩码的掩码来掩盖网络的不同参数。我该怎么办？

换句话说，我只想训练非零参数

nn.Linear(input_dim,output_dim)*W

并保持所有其他参数等于零。