Don't enforce docstrings on bool dispatch (#15306)

Summary:
Allows 2 functions that are boolean dispatched to have no docstrings (the only case that will fail now is if both functions have docstrings)

Fixes #15281
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15306

Differential Revision: D13494884

Pulled By: driazati

fbshipit-source-id: 65fec39ae03a7d6a68ad617c9b270faeb1617930

diff --git a/torch/_jit_internal.py b/torch/_jit_internal.py
index 18efc48..5594d41 100644 (file)
--- a/torch/_jit_internal.py
+++ b/torch/_jit_internal.py
@@ -137,6 +137,9 @@ def boolean_dispatch(arg_name, arg_index, default, if_true, if_false):
     elif if_false.__doc__ is None and if_true.__doc__ is not None:
         doc = if_true.__doc__
         if_false.__doc__ = doc
+    elif if_false.__doc__ is None and if_true.__doc__ is None:
+        # neither function has a docstring
+        doc = None
     else:
         raise RuntimeError("only one function can have a docstring")
     fn.__doc__ = doc