imt-ring 1.3.6__py3-none-any.whl → 1.3.7__py3-none-any.whl

{imt_ring-1.3.6.dist-info → imt_ring-1.3.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.6
+Version: 1.3.7
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.3.6.dist-info → imt_ring-1.3.7.dist-info}/RECORD

@@ -56,7 +56,7 @@ ring/ml/callbacks.py,sha256=DkSy5c7IRqAAks2dx8acEBExYxUv-xiUFwZn4odPYq4,13253
 ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=OWRDu2COmptzbpJWlRLbPIn_ioKZCAd_iu-eiY_aPjk,8521
-ring/ml/train.py,sha256=ftt2MOSSNGCdL7ZoAXcbIgeHW1Wkpgp6XYyLIBUIClI,10872
+ring/ml/train.py,sha256=uDW6JMdbMcjUKr3wCL2drWzDUd0Pc3BoroUwLcYoUx4,10914
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
@@ -77,7 +77,7 @@ ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
 ring/utils/utils.py,sha256=mIcKNv5v2de8HrG7bAhl2bNfmwkMZyIIwFkJq2XWMOI,5357
-imt_ring-1.3.6.dist-info/METADATA,sha256=E5mVtL-2o6-U-Ov56yd4M0RVQs0VJoLSezHpBWGtleg,3104
-imt_ring-1.3.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.6.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.6.dist-info/RECORD,,
+imt_ring-1.3.7.dist-info/METADATA,sha256=V6Oow_ZZwpBuHuIbyPIoKFtrhFboxMmuIPx1Rilq3-A,3104
+imt_ring-1.3.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.3.7.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.3.7.dist-info/RECORD,,

ring/ml/train.py

@@ -5,6 +5,8 @@ from typing import Callable, Optional, Tuple
 import jax
 import jax.numpy as jnp
 import optax
+import tree_utils
+
 from ring import maths
 from ring.algorithms.generator import types
 from ring.ml import base as ml_base
@@ -15,8 +17,6 @@ from ring.utils import distribute_batchsize
 from ring.utils import expand_batchsize
 from ring.utils import parse_path
 from ring.utils import pickle_load
-import tree_utils
-
 import wandb
 
 # (T, N, F) -> Scalar
@@ -142,15 +142,17 @@ def train_fn(
         Wether or not the training run was killed by a callback.
     """
 
+    filter = filter.nojit()
+
     if checkpoint is not None:
         checkpoint = Path(checkpoint).with_suffix(".pickle")
         recv_checkpoint: dict = pickle_load(checkpoint)
-        filter.params = recv_checkpoint["params"]
+        filter_params = recv_checkpoint["params"]
         opt_state = recv_checkpoint["opt_state"]
+        del recv_checkpoint
+    else:
+        filter_params = filter.search_attr("params")
 
-    filter = filter.nojit()
-
-    filter_params = filter.search_attr("params")
     if filter_params is None:
         X, _ = generator(jax.random.PRNGKey(1))
         filter_params, _ = filter.init(X=X, seed=seed_network)