Making onnx-safe option overrideable

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

megatron/arguments.py

@@ -169,7 +169,7 @@ def _add_network_size_args(parser):
                        help='Use OpenAIs GeLU implementation. This option'
                        'should not be used unless for backward compatibility'
                        'reasons.')
-    group.add_argument('--onnx-safe', action='store_true',
+    group.add_argument('--onnx-safe', type=bool, required=False,
                        help='Use workarounds for known problems with Torch ONNX exporter')
 
     return parser