How to initialize weights with one line

With mynet = nn.Sequential(nn.Linear(1,10), nn.Sigmoid()), how to initialize weights with one line as ones or zeros instead of as the default random weights?

def init_weights(m):
    if isinstance(m, nn.Linear):
        m.weight.data.fill_(1)
        m.bias.data.fill_(0)

net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)

For more information, check this Stackoverflow answer

1 Like