I think this can be done manually only by zeroing your weights after loss.backward().
Check this topic:

  
    
    
    Freezing part of the Layer Weights 
  
  
    Hello Everyone, 
How could I freeze some parts of the layer weights to zero and not the entire layer. 
I tried below code, but it doesn’t freeze the specific parts(1:10 array in 2nd dimension) of the layer weights. 
I am new to ML & started with Pytorch. Appreciate any help. Thanks. 
for child in model_ft.children(): 
print(“Freezing Parameters(1->10) on the Convolution Layer”,child) 
for param in child.parameters(): 
param.data[:,1:10,:,:].zero_() 
param.data[:,1:10,:,:].requires_grad = False 
o…