Does TorchInfo consider parameters of a layer that went through torch.no_grad() as trainable ?

[maelleF]

Hi !
I have a flexible model that I wanted to finetuned at different levels of layers, so for that I added the following code in my forward function:

with torch.no_grad():
    for net in self.layers[:train_start_index]:
        x = net(x)
for net in self.layers[train_start_index:output_index+1]:
    x = net(x) `

Even if the model works as intended, while running the summary function of torch info, I always have the same amount of trainable parameters. Is there something I missed in the use of torchinfo, or is it just not something currently implemented ?

Thanks for your help !