Thanks for your suggestion. But when I apply these changes, it still doesn’t train… Here is my updated code:
def init_weights(m):
  with torch.no_grad():
    if type(m) == nn.Linear or type(m) == nn.Conv2d:
      m.weight = nn.Parameter(custom_weight_fn(m.weight.shape))
      m.bias = nn.Parameter(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) #if I comment this line, the model trains with no problem

I also tried m.param.copy_ and unfortunately that doesn’t work either.