fluidflow 0.2.0__tar.gz → 0.2.2__tar.gz

{fluidflow-0.2.0 → fluidflow-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fluidflow
-Version: 0.2.0
+Version: 0.2.2
 Summary: FluidFlow: models and training utilities for flow matching and diffusion
 Author-email: David Ramos <david.ramos.archilla@upm.es>
 Requires-Python: >=3.8

{fluidflow-0.2.0 → fluidflow-0.2.2}/fluidFlow/trainer.py

@@ -10,7 +10,6 @@ from torch.profiler import profile, ProfilerActivity, schedule
 
 from accelerate import Accelerator, DataLoaderConfiguration
 from accelerate import FullyShardedDataParallelPlugin
-from torch.distributed.fsdp import ShardingStrategy
 from torch.distributed.fsdp import MixedPrecisionPolicy
 from ema_pytorch import EMA
 
@@ -142,7 +141,7 @@ class Trainer(object):
             # Separate parameters for Muon (2D) and AdamW (1D)
             neural_net = diffusion_model.neural_net if hasattr(diffusion_model, "neural_net") else diffusion_model
             muon_params = neural_net.get_2d_params()
-            adam_params = neural_net.get_1d_params()
+            adam_params = neural_net.get_not_2d_params()
             self.opt_muon = torch.optim.Muon(muon_params, lr=train_lr, weight_decay=1e-3)
             self.opt_adam = AdamW(adam_params, lr=train_lr, betas=adam_betas, weight_decay=1e-2, fused=True, eps=eps)
             self.opts = [self.opt_muon, self.opt_adam]
@@ -442,4 +441,4 @@ class Trainer(object):
         plt.title('Training Loss Evolution')
         plt.legend()
         plt.savefig(self.results_folder / "loss_evolution.png", bbox_inches="tight", pad_inches=0)
-        plt.close()
+        plt.close()

{fluidflow-0.2.0 → fluidflow-0.2.2}/fluidflow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fluidflow
-Version: 0.2.0
+Version: 0.2.2
 Summary: FluidFlow: models and training utilities for flow matching and diffusion
 Author-email: David Ramos <david.ramos.archilla@upm.es>
 Requires-Python: >=3.8

{fluidflow-0.2.0 → fluidflow-0.2.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "fluidflow"
 authors = [{ name = "David Ramos", email = "david.ramos.archilla@upm.es" }]
-version = "0.2.0"
+version = "0.2.2"
 description = "FluidFlow: models and training utilities for flow matching and diffusion"
 readme = "README.md"
 requires-python = ">=3.8"