Order torch imports after Julia init

test/test_torch.py

@@ -2,7 +2,6 @@ import unittest
 import numpy as np
 import pandas as pd
 from pysr import sympy2torch, PySRRegressor
-import torch
 import sympy
 from functools import partial
 
@@ -14,6 +13,8 @@ class TestTorch(unittest.TestCase):
     def test_sympy2torch(self):
         x, y, z = sympy.symbols("x y z")
         cosx = 1.0 * sympy.cos(x) + y
+
+        import torch
         X = torch.tensor(np.random.randn(1000, 3))
         true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
         torch_module = sympy2torch(cosx, [x, y, z])
@@ -49,6 +50,8 @@ class TestTorch(unittest.TestCase):
 
         tformat = model.pytorch()
         self.assertEqual(str(tformat), "_SingleSymPyModule(expression=cos(x1)**2)")
+        import torch
+
         np.testing.assert_almost_equal(
             tformat(torch.tensor(X.values)).detach().numpy(),
             np.square(np.cos(X.values[:, 1])),  # Selection 1st feature
@@ -81,6 +84,8 @@ class TestTorch(unittest.TestCase):
 
         tformat = model.pytorch()
         self.assertEqual(str(tformat), "_SingleSymPyModule(expression=cos(x1)**2)")
+
+        import torch
         np.testing.assert_almost_equal(
             tformat(torch.tensor(X)).detach().numpy(),
             np.square(np.cos(X[:, 1])),  # 2nd feature
@@ -93,6 +98,7 @@ class TestTorch(unittest.TestCase):
 
         module = sympy2torch(expression, [x, y, z])
 
+        import torch
         X = torch.rand(100, 3).float() * 10
 
         true_out = (
@@ -133,6 +139,7 @@ class TestTorch(unittest.TestCase):
 
         tformat = model.pytorch()
         self.assertEqual(str(tformat), "_SingleSymPyModule(expression=sin(x1))")
+        import torch
         np.testing.assert_almost_equal(
             tformat(torch.tensor(X)).detach().numpy(),
             np.sin(X[:, 1]),
@@ -152,6 +159,7 @@ class TestTorch(unittest.TestCase):
         torch_module = model.pytorch()
 
         np_output = model.predict(X.values)
+        import torch
         torch_output = torch_module(torch.tensor(X.values)).detach().numpy()
 
         np.testing.assert_almost_equal(np_output, torch_output, decimal=4)