do you mean something like this,
import torch.nn as nn, torch

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.A = nn.ModuleDict([['lin', nn.Linear(3, 3)], 
                             ['dropout', nn.Dropout(0.5)]])
    self.B = nn.Linear(3, 3)
  def forward(self, x):
    out = self.B(x)
    out = self.A['lin'](out)
    out = self.A['dropout'](out)
    return out

net = Net()

optimizer1 = torch.optim.SGD(net.A.parameters(), lr=0.01)
optimizer2 = torch.optim.SGD(net.B.parameters(), lr=0.01)

for i in range(2):
  optimizer1.zero_grad()
  optimizer2.zero_grad()
  input = torch.randn(3, 3)
  loss = (net(input)).sum()
  loss.backward()
  if loss < 1: # replace with net(input) == prediction
    optimizer1.step()
  else:
    optimizer2.step()

or with one optimizer
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)

for i in range(2):
  optimizer.zero_grad()
  input = torch.randn(3, 3)
  loss = (net(input)).sum()
  print(loss)
  loss.backward()
  if loss < 1: # replace with net(input) == prediction
    net.A.lin.requires_grad_(False)
    net.B.requires_grad_(True)
  else:
    net.A.lin.requires_grad_(True)
    net.B.requires_grad_(False)
  print(list(net.parameters()))
  optimizer.step()