model = torchvision.models.resnet18(pretrained=False)

class SimpleCustomLayer(nn.Module):
    def __init__(self):
        super(SimpleCustomLayer).__init__()

    def forward(self, x):
        return x


block = nn.Sequential(SimpleCustomLayer(), nn.ReLU())
for i, l in enumerate(model.modules()):
    if isinstance(l, nn.ReLU):
        print(f"Replacing at {i}")
        model[i] = block


This gives the following error which is very understandable.
Traceback (most recent call last):
  File "/path/", line 23, in <module>
    model[i] = block
TypeError: 'ResNet' object does not support item assignment

Does PyTorch provide any API for switching such layers ?