Add docstring for sympy2torch

pysr/export_torch.py

@@ -185,8 +185,11 @@ def _initialize_torch():
 
 
 def sympy2torch(expression, symbols_in):
-    global torch
-    global _Node
+    """Returns a module for a given sympy expression with trainable parameters;
+
+    This function will assume the input to the module is a matrix X, where
+        each column corresponds to each symbol you pass in `symbols_in`.
+    """
     global SingleSymPyModule
 
     _initialize_torch()