TL;DRIf you create weight implicitly by creating a linear layer, you should set modle='fan_in'. linear = torch.nn.Linear(node_in, node_out) init.kaiming_normal_(linear.weight, mode=’fan_in’) t = relu(linear(x_valid))If you create weight explicitly by creating a random matrix, you should set modle='fan_out'. w1 = torch.randn(node_in, node_out) init.kaiming_normal_(w1, mode=’fan_out’) b1 = torch.ran