I have a model that trains fine without weight initialization but when I try to initialize the weights of my model to custom values (generated by custom_weight_fn), it stops training. Here is the code:
def init_weights(m):
  if type(m) == nn.Linear or type(m) == nn.Conv2d:
    m.weight.data=custom_weight_fn(m.weight.shape)
    m.bias.data=custom_weight_fn(m.bias.shape)

model = nn.Sequential(
    nn.Conv2d(3, channel_1, 5, padding=2),
    nn.ReLU(),
    nn.Conv2d(channel_1, channel_2, 3, padding=1),
    nn.ReLU(),
    Flatten(),
    nn.Linear(32*32*channel_2, 10)
)

model.apply(init_weights)


Does anybody know why this is not working? The code runs without any errors but the model doesn’t train.
Thanks!