MilesCranmer commited on
Commit
5408936
1 Parent(s): a5d2fe9

Update torch export docs

Browse files
Files changed (1) hide show
  1. pysr/export_torch.py +1 -7
pysr/export_torch.py CHANGED
@@ -1,8 +1,3 @@
1
- #####
2
- # From https://github.com/patrick-kidger/sympytorch
3
- # Copied here to allow PySR-specific tweaks
4
- #####
5
-
6
  import collections as co
7
  import functools as ft
8
  import sympy
@@ -23,7 +18,7 @@ def _initialize_torch():
23
  global sympytorch
24
  global PySRTorchModule
25
 
26
- # Way to lazy load torch, only if this is called,
27
  # but still allow this module to be loaded in __init__
28
  if not torch_initialized:
29
  try:
@@ -35,7 +30,6 @@ def _initialize_torch():
35
 
36
 
37
  class PySRTorchModule(torch.nn.Module):
38
- """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""
39
  def __init__(self, *, expression, symbols_in,
40
  selection=None, extra_funcs=None, **kwargs):
41
  super().__init__(**kwargs)
 
 
 
 
 
 
1
  import collections as co
2
  import functools as ft
3
  import sympy
 
18
  global sympytorch
19
  global PySRTorchModule
20
 
21
+ # Way to lazy load torch and sympytorch, only if this is called,
22
  # but still allow this module to be loaded in __init__
23
  if not torch_initialized:
24
  try:
 
30
 
31
 
32
  class PySRTorchModule(torch.nn.Module):
 
33
  def __init__(self, *, expression, symbols_in,
34
  selection=None, extra_funcs=None, **kwargs):
35
  super().__init__(**kwargs)