Can I use model.apply with argument in function?

1

Something like

def set_dropout(m, p):
    if isinstance(m, nn.Dropout2d):
        m.p = p
netG.apply(set_dropout(0.1))
2

Yes, you can use a lambda for this:

def set_dropout(m, p):
    if isinstance(m, nn.Dropout):
        m.p = p

model.apply(lambda m: set_dropout(m, 0.1))

Note that you are checking for nn.Dopout2d, while the “vanilla dropout” is nn.Dropout.
If you are using nn.Dropout2d in your model, please ignore this sentence. :wink:

1 Like