I am not sure if any predefined way exists in pytorch.
You may need to write a custom linear layer to achieve it, something like below:

  

      gist.github.com
  

  
    https://gist.github.com/InnovArul/ee2986feff215fed94401ec83367408b 3

  linear_partial_freeze.py
  import torch, torch.nn as nn                                                         
import torch.optim as optim, torch.nn.functional as F                                
                                                                                     
class CustomLinear(nn.Module):                                                       
    def __init__(self, mask):                                                        
        super().__init__()                                                           
        self.register_buffer("mask", mask)                                           
        out_channels, in_channels = mask.shape                                       
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels))       
        
    This file has been truncated. show original