fix(ver): consider torch version for function_call import

torchimize/optimizer/gna_opt.py

@@ -4,6 +4,16 @@
 from torch.optim.optimizer import Optimizer
 from typing import List
 
+from distutils.version import LooseVersion
+if LooseVersion(torch.__version__) > LooseVersion('2.0.0'):
+    from torch.func import functional_call
+else:
+    try:
+        from torch.nn.utils.stateless import functional_call
+    except ImportError:
+        from torch.nn.utils._stateless import functional_call
+
+
 
 class GNA(Optimizer):
     r"""Implements Gauss-Newton.
@@ -62,7 +72,7 @@ def step(self, x: torch.Tensor, closure=None):
             if self.hessian_approx:
                 # vectorized jacobian (https://github.com/pytorch/pytorch/issues/49171)
                 def func(*params: torch.Tensor):
-                    out = torch.func.functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
+                    out = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
                     return out
                 self._j_list: tuple[torch.Tensor] = torch.autograd.functional.jacobian(func, values, create_graph=False)    # NxCxBxCxHxW
                 # create hessian approximation
@@ -77,7 +87,7 @@ def func(*params: torch.Tensor):
             else:
                 # vectorized hessian (https://github.com/pytorch/pytorch/issues/49171)
                 def func(*params: torch.Tensor):
-                    out: torch.Tensor = torch.func.functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
+                    out: torch.Tensor = functional_call(self._model, {n: p for n, p in zip(keys, params)}, x)
                     return out.square().sum()
                 self._h_list: tuple[torch.Tensor] = torch.autograd.functional.hessian(func, tuple(self._model.parameters()), create_graph=False)
                 self._h_list = [self._h_list[i][i] for i in range(len(self._h_list))] # filter j-th element