1 files changed, 5 insertions, 5 deletions

diff --git a/data/csv.py b/data/csv.py
index 55a1988..d9f9db8 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -118,7 +118,7 @@ class VlpnDataModule():
         size: int = 768,
         num_buckets: int = 0,
         bucket_step_size: int = 64,
-        max_pixels_per_bucket: Optional[int] = None,
+        bucket_max_pixels: Optional[int] = None,
         progressive_buckets: bool = False,
         dropout: float = 0,
         interpolation: str = "bicubic",
@@ -145,7 +145,7 @@ class VlpnDataModule():
         self.size = size
         self.num_buckets = num_buckets
         self.bucket_step_size = bucket_step_size
-        self.max_pixels_per_bucket = max_pixels_per_bucket
+        self.bucket_max_pixels = bucket_max_pixels
         self.progressive_buckets = progressive_buckets
         self.dropout = dropout
         self.template_key = template_key
@@ -234,7 +234,7 @@ class VlpnDataModule():
         train_dataset = VlpnDataset(
             self.data_train, self.prompt_processor,
             num_buckets=self.num_buckets, progressive_buckets=self.progressive_buckets,
-            bucket_step_size=self.bucket_step_size, max_pixels_per_bucket=self.max_pixels_per_bucket,
+            bucket_step_size=self.bucket_step_size, bucket_max_pixels=self.bucket_max_pixels,
             batch_size=self.batch_size, generator=generator,
             size=self.size, interpolation=self.interpolation,
             num_class_images=self.num_class_images, dropout=self.dropout, shuffle=True,
@@ -264,7 +264,7 @@ class VlpnDataset(IterableDataset):
         prompt_processor: PromptProcessor,
         num_buckets: int = 1,
         bucket_step_size: int = 64,
-        max_pixels_per_bucket: Optional[int] = None,
+        bucket_max_pixels: Optional[int] = None,
         progressive_buckets: bool = False,
         batch_size: int = 1,
         num_class_images: int = 0,
@@ -290,7 +290,7 @@ class VlpnDataset(IterableDataset):
             base_size=size,
             step_size=bucket_step_size,
             num_buckets=num_buckets,
-            max_pixels=max_pixels_per_bucket,
+            max_pixels=bucket_max_pixels,
             progressive_buckets=progressive_buckets,
         )