[mednet] Fixes after rebase

src/mednet/libs/classification/scripts/train.py

@@ -76,7 +76,6 @@ def train(
         datamodule,
         model,
         batch_size,
-        batch_chunk_count,
         drop_incomplete_batch,
         cache_samples,
         parallel,
@@ -120,6 +119,6 @@ def train(
         max_epochs=epochs,
         output_folder=output_folder,
         monitoring_interval=monitoring_interval,
-        batch_chunk_count=batch_chunk_count,
+        accumulate_grad_batches=accumulate_grad_batches,
         checkpoint=checkpoint_file,
     )

src/mednet/libs/common/scripts/predict.py

@@ -127,7 +127,6 @@ def setup_datamodule(
     parallel,
 ) -> None:  # numpydoc ignore=PR01
     """Configure and set up the datamodule."""
-    datamodule.set_chunk_size(batch_size, 1)
     datamodule.parallel = parallel
     datamodule.model_transforms = model.model_transforms
 

src/mednet/libs/common/scripts/train.py

@@ -276,13 +276,11 @@ def setup_datamodule(
     datamodule,
     model,
     batch_size,
-    batch_chunk_count,
     drop_incomplete_batch,
     cache_samples,
     parallel,
 ) -> None:  # numpydoc ignore=PR01
     """Configure and set up the datamodule."""
-    datamodule.set_chunk_size(batch_size, batch_chunk_count)
     datamodule.drop_incomplete_batch = drop_incomplete_batch
     datamodule.cache_samples = cache_samples
     datamodule.parallel = parallel

src/mednet/libs/segmentation/scripts/experiment.py

@@ -41,7 +41,7 @@ def experiment(
     output_folder,
     epochs,
     batch_size,
-    batch_chunk_count,
+    accumulate_grad_batches,
     drop_incomplete_batch,
     datamodule,
     validation_period,
@@ -79,7 +79,7 @@ def experiment(
         output_folder=train_output_folder,
         epochs=epochs,
         batch_size=batch_size,
-        batch_chunk_count=batch_chunk_count,
+        accumulate_grad_batches=accumulate_grad_batches,
         drop_incomplete_batch=drop_incomplete_batch,
         datamodule=datamodule,
         validation_period=validation_period,

src/mednet/libs/segmentation/scripts/train.py

@@ -32,7 +32,7 @@ def train(
     output_folder,
     epochs,
     batch_size,
-    batch_chunk_count,
+    accumulate_grad_batches,
     drop_incomplete_batch,
     datamodule,
     validation_period,
@@ -62,7 +62,6 @@ def train(
         datamodule,
         model,
         batch_size,
-        batch_chunk_count,
         drop_incomplete_batch,
         cache_samples,
         parallel,
@@ -81,7 +80,7 @@ def train(
         device_manager,
         epochs,
         batch_size,
-        batch_chunk_count,
+        accumulate_grad_batches,
         drop_incomplete_batch,
         validation_period,
         cache_samples,
@@ -98,6 +97,6 @@ def train(
         max_epochs=epochs,
         output_folder=output_folder,
         monitoring_interval=monitoring_interval,
-        batch_chunk_count=batch_chunk_count,
+        accumulate_grad_batches=accumulate_grad_batches,
         checkpoint=checkpoint_file,
     )