RRAEsTorch 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

RRAEsTorch/AE_base/AE_base.py

@@ -1,5 +1,4 @@
 from RRAEsTorch.utilities import MLP_with_linear 
-import jax.random as jrandom
 import torch.nn as nn
 from torch.func import vmap
 

RRAEsTorch/training_classes/training_classes.py

@@ -25,7 +25,6 @@ import matplotlib.pyplot as plt
 from prettytable import PrettyTable
 import torch
 from torch.utils.data import TensorDataset, DataLoader
-from adabelief_pytorch import AdaBelief
 
 class Circular_list:
     """
@@ -221,7 +220,7 @@ class Trainor_class:
         tracker=Null_Tracker(),
         stagn_window=20,
         eps_fn=lambda lat, bs: None,
-        optimizer=AdaBelief,
+        optimizer=torch.optim.Adam,
         verbatim = {
                     "print_type": "std",
                     "window_size" : 5,  
@@ -307,7 +306,7 @@ class Trainor_class:
         for steps, lr, batch_size in zip(step_st, lr_st, batch_size_st):
             try:
                 t_t = 0.0  # Zero time
-                optimizer = torch.optim.Adam(
+                optimizer_tr = optimizer(
                     filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3
                 )
 
@@ -341,11 +340,11 @@ class Trainor_class:
                     step_kwargs = merge_dicts(loss_kwargs, track_params)
 
                     # Compute loss
-                    loss, model, optimizer, (aux, extra_track) = make_step(
+                    loss, model, optimizer_tr, (aux, extra_track) = make_step(
                                                                 model,
                                                                 input_b,
                                                                 out_b,
-                                                                optimizer,
+                                                                optimizer_tr,
                                                                 idx_b,
                                                                 epsilon,
                                                                 **step_kwargs,
@@ -391,7 +390,7 @@ class Trainor_class:
                         self.del_file(f"checkpoint_k_{track_params.get('k_max')}")
                         model = self.model
 
-                        optimizer = torch.optim.Adam(
+                        optimizer_tr = optimizer(
                             filter(lambda p: p.requires_grad, model.parameters()), lr=lr
                         )
 
@@ -974,4 +973,4 @@ class RRAE_Trainor_class(AE_Trainor_class):
             latent_func=call_func,
             decode_func=decode_func,
             norm_out_func=norm_out_func,
-        )
+        )

{rraestorch-0.1.0.dist-info → rraestorch-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: RRAEsTorch
-Version: 0.1.0
+Version: 0.1.2
 Summary:  A repo for RRAEs in PyTorch.
 Author-email: Jad Mounayer <jad.mounayer@outlook.com>
 License: MIT

{rraestorch-0.1.0.dist-info → rraestorch-0.1.2.dist-info}/RECORD

@@ -1,6 +1,6 @@
 RRAEsTorch/__init__.py,sha256=f234R6usRCqIgmBmiXyZNIHa7VrDe5E-KZO0Y6Ek5AQ,33
 RRAEsTorch/config.py,sha256=bQPwc_2KTvhglH_WIRSb5_6CpUQQj9AGpfqBp8_kuys,2931
-RRAEsTorch/AE_base/AE_base.py,sha256=4nnuv3VbrakkvpmcT9NqRBuFUZsC1BGMpQWdpxEcGyk,3398
+RRAEsTorch/AE_base/AE_base.py,sha256=Eeo_I7p5P-357rnOmCuFxosJgmBg4KPyMA8n70sTV7U,3368
 RRAEsTorch/AE_base/__init__.py,sha256=95YfMgEWzIFAkm--Ci-a9YPSGfCs2PDAK2sbfScT7oo,24
 RRAEsTorch/AE_classes/AE_classes.py,sha256=LTIKobJ5FXOPwGZLRhs82NKeuQzvXzj2YzdhuBLi2i0,18108
 RRAEsTorch/AE_classes/__init__.py,sha256=inM2_YPJG8T-lwx-CUg-zL2EMltmROQAlNZeZmnvVGA,27
@@ -15,13 +15,13 @@ RRAEsTorch/tests/test_wrappers.py,sha256=Ike4IfMUx2Qic3f3_cBikgFPEU1WW5TuH1jT_r2
 RRAEsTorch/trackers/__init__.py,sha256=3c9qcUMZiUfVr93rxFp6l11lIDthyK3PCY_-P-sNX3I,25
 RRAEsTorch/trackers/trackers.py,sha256=Pn1ejMxMjAtvgDazFFwa3qiZhogG5GtXj4UIIFiBpuY,9127
 RRAEsTorch/training_classes/__init__.py,sha256=K_Id4yhw640jp2JN15-0E4wJi4sPadi1fFRgovMV3kw,101
-RRAEsTorch/training_classes/training_classes.py,sha256=l1wPWZghYwzz_wxgrBi2fmOumgjbCz3qaaIcDks-GCg,36668
+RRAEsTorch/training_classes/training_classes.py,sha256=a7JjhCrH7s7VmVCsvKj768Ciq-tdbh6E_B9aG1kw7vc,36634
 RRAEsTorch/utilities/__init__.py,sha256=NtlizCcRW4qcsULXxWfjPk265rLJst0-GqWLRah2yDY,26
 RRAEsTorch/utilities/utilities.py,sha256=FzJWV9oFPF9sL9MC2m7euMqMKxCuLUEukzLfU0cF2to,53396
 RRAEsTorch/wrappers/__init__.py,sha256=txiLh4ylnuvPlapagz7DiAslmjllOzTqwCDL2dFr6dM,44
 RRAEsTorch/wrappers/wrappers.py,sha256=9Rmq2RS_EkZvsg96SKrt1HFIP35sF0xyPI0goV0ujOs,9659
-rraestorch-0.1.0.dist-info/METADATA,sha256=Wm_NiD_XGPLm8bHVI-U2Rt7aDnNAQ1tkoeeo3pzWarQ,3055
-rraestorch-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rraestorch-0.1.0.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
-rraestorch-0.1.0.dist-info/licenses/LICENSE copy,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
-rraestorch-0.1.0.dist-info/RECORD,,
+rraestorch-0.1.2.dist-info/METADATA,sha256=Gf8TImFT4nLlgXlaHjJPvzlj75Q7nBVrOT29eanTubA,3055
+rraestorch-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rraestorch-0.1.2.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.2.dist-info/licenses/LICENSE copy,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.2.dist-info/RECORD,,