How to split neural network parameters into groups

Question

I have the following neural network toy example:

class Neural_net(nn.Module):
    def __init__(self):
        super(Neural_net, self).__init__()
        self.fc1    = nn.Linear(2, 2)        
        self.fc2    = nn.Linear(2, 1)
        self.fc_out = nn.Linear(1, 1)      
        
    def forward(self, x,train = True):
        x = torch.tanh(self.fc1(x))
        x = torch.tanh(self.fc2(x))
        x = self.fc_out(x)
        return x

net = Neural_net()

I also have a dictionary of the same structure as net.named_parameters(), whose entries are either True or False. So each parameter in the network has a corresponding flag marking it as True or False. The dictionary is created based on the answer by @ibarrond to this question: How to iterate through all parameters in a neural network using PyTorch.

How can I use this dictionary to split the network parameters into two groups? I want to place the parameters tagged as True in one group and those with a False tag in another group.

I was thinking something of the form:

    parameter_group_1 = [{"params": [p for n, p in net.named_parameters() if cond_parameters[n]==True], 'lr': h1}]

but this returns the following error

Boolean value of Tensor with more than one value is ambiguous

Since the value cond_parameters[n] corresponding to the key n has more than one entry.

So the problem is that when doing

for n, p in net.named_parameters():

the returned parameter p is not a single value, but rather a tensor containing e.g the weights for the first layer, the biases for the first layer, etc. And similarly, cond_parameters[n] is not a single value but a tensor with True or False entries of the same size as the corresponding parameter tensor p.