Adjust init

Author: JasonGross

gbmi/exp_modular_arithmetic/train.py

@@ -192,6 +192,7 @@ def log_softmax(self, x: Tensor, **kwargs) -> Tensor:
         p=113,
         training_ratio=0.8,
         logging_options=ModelMatrixLoggingOptions.none(),
+        init_mode="kaiming_uniform",
     ),
     seed=0,
     deterministic=False,