fix batch_size bug

Author: FilippoOlivo

pina/_src/data/creator.py

@@ -79,12 +79,14 @@ def _compute_batch_sizes(self, datasets):
         """
         batch_sizes = {}
         if self.batching_mode == "common_batch_size":
-            
+
             if self.batch_size is None:
-                batch_size = max(dataset.length for dataset in datasets.values())
+                batch_size = max(
+                    dataset.length for dataset in datasets.values()
+                )
             else:
                 batch_size = self.batch_size
-            
+
             for name in datasets.keys():
                 batch_sizes[name] = min(batch_size, len(datasets[name]))
             return batch_sizes
@@ -169,9 +171,12 @@ def __call__(self, datasets):
         dataloaders = {}
         if self.batching_mode == "common_batch_size":
             max_len = max(len(dataset) for dataset in datasets.values())
-
+        print(batch_sizes)
         for name, dataset in datasets.items():
-            if self.batching_mode == "common_batch_size" and dataset.length != batch_sizes[name]:
+            if (
+                self.batching_mode == "common_batch_size"
+                and dataset.length != batch_sizes[name]
+            ):
                 dataset.max_len = max_len
             dataloaders[name] = self.conditions[name].create_dataloader(
                 dataset=dataset,

tests/conftest.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 import pytest
 
+
 @pytest.fixture
 def clean_tmp_dir(tmp_path):
     path = Path(tmp_path)
@@ -13,4 +14,4 @@ def clean_tmp_dir(tmp_path):
     yield path
 
     if path.exists():
-        shutil.rmtree(path)
+        shutil.rmtree(path)

tests/test_problem.py

@@ -48,6 +48,7 @@ def test_variables_correct_order_sampling():
         poisson_problem.input_variables
     )
 
+
 def test_add_points():
     poisson_problem = Poisson()
     poisson_problem.discretise_domain(1, "random", domains=["D"])
@@ -90,4 +91,4 @@ def test_wrong_custom_sampling_logic(mode):
 
     # Necessary cleanup
     if "new" in poisson_problem.domains:
-        del poisson_problem.domains["new"]
+        del poisson_problem.domains["new"]

tests/test_solver/test_competitive_pinn.py

@@ -149,4 +149,4 @@ def test_train_load_restore(clean_tmp_dir, problem):
     )
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
+    )

tests/test_solver/test_pinn.py

@@ -116,6 +116,7 @@ def test_train_load_restore(clean_tmp_dir, problem):
     )
     trainer.train()
     import os
+
     print(os.listdir(f"{dir}/lightning_logs/version_0/checkpoints/"))
 
     # restore
@@ -125,7 +126,6 @@ def test_train_load_restore(clean_tmp_dir, problem):
         + "epoch=4-step=5.ckpt"
     )
 
-
     # loading
     new_solver = PINN.load_from_checkpoint(
         f"{dir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt",
@@ -138,4 +138,4 @@ def test_train_load_restore(clean_tmp_dir, problem):
     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
+    )

tests/test_solver/test_rba_pinn.py

@@ -158,4 +158,4 @@ def test_train_load_restore(clean_tmp_dir, problem):
     )
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
+    )

tests/test_solver/test_supervised_solver.py

@@ -216,6 +216,7 @@ def test_solver_test_graph(batch_size, use_lt):
 from pathlib import Path
 import pytest
 
+
 @pytest.fixture
 def clean_tmp_dir():
     path = Path("tests/test_solver/tmp/")
@@ -229,6 +230,7 @@ def clean_tmp_dir():
     if path.exists():
         shutil.rmtree(path)
 
+
 def test_train_load_restore(clean_tmp_dir):
     dir = clean_tmp_dir
     problem = LabelTensorProblem()
@@ -264,4 +266,4 @@ def test_train_load_restore(clean_tmp_dir):
     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
+    )