would this work,
class Model(nn.Module):
  def __init__(self):
    super().__init__()
    self.module_list = nn.ModuleList([nn.Linear(2, 2, bias=False) for i in range(5)])
  def forward(self, input, num_layers_to_drop):
    x = torch.randint(0, len(self.module_list), (num_layers_to_drop,))
    # x is index for layers we want to drop
    print(x)
    out = input
    for index, layer in enumerate(self.module_list):
      if index not in x:
        out = layer(out) # we pass input only through indices not in x
    return out

net = Model()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
input = torch.randn(2, 2)
optimizer.zero_grad()
loss = net(input, 2).sum()
loss.backward()

before optimizer.step()
list(net.module_list.parameters())

[Parameter containing:
 tensor([[ 0.2603, -0.4476],
         [ 0.2753, -0.3421]], requires_grad=True), Parameter containing:
 tensor([[ 0.1724, -0.2733],
         [-0.0699,  0.0609]], requires_grad=True), Parameter containing:
 tensor([[ 0.7007,  0.4430],
         [ 0.1900, -0.0225]], requires_grad=True), Parameter containing:
 tensor([[0.1543, 0.5798],
         [0.2771, 0.6099]], requires_grad=True), Parameter containing:
 tensor([[ 0.3549, -0.3747],
         [ 0.1592, -0.2938]], requires_grad=True)]

we do our update, suppose x is tensor([2, 0])
optimizer.step()

[Parameter containing:
 tensor([[ 0.2603, -0.4476],
         [ 0.2753, -0.3421]], requires_grad=True), Parameter containing:
 tensor([[ 0.1715, -0.2729],
         [-0.0708,  0.0613]], requires_grad=True), Parameter containing:
 tensor([[ 0.7007,  0.4430],
         [ 0.1900, -0.0225]], requires_grad=True), Parameter containing:
 tensor([[0.1556, 0.5793],
         [0.2754, 0.6105]], requires_grad=True), Parameter containing:
 tensor([[ 0.3549, -0.3745],
         [ 0.1591, -0.2937]], requires_grad=True)]

layer at index 0, 2 do not get updated.