Error Freezing Weights

[mostafaelhoushi]

I am trying to freeze weights of linear layers but I get this error:

module.weight.requires_grad = not freeze
RuntimeError: you can only change requires_grad flags of leaf variables. If you want to use a computed variable in a subgraph that doesn't require differentiation use var_no_grad = var.detach().

and in my case the module is ColumnParallelLinear whose weight does seem to be a leaf variable:

fairscale/fairscale/nn/model_parallel/layers.py

Line 262 in 164cc0f

self.weight = Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))

so I am not sure why I am getting this error message.