Improved sample output and progress bars

1 files changed, 5 insertions, 5 deletions

diff --git a/data/dreambooth/csv.py b/data/dreambooth/csv.py
index 14c13bb..85ed4a5 100644
--- a/data/dreambooth/csv.py
+++ b/data/dreambooth/csv.py
@@ -108,14 +108,14 @@ class CSVDataset(Dataset):
         else:
             self.class_data_root = None
 
-        self.interpolation = {"linear": PIL.Image.LINEAR,
-                              "bilinear": PIL.Image.BILINEAR,
-                              "bicubic": PIL.Image.BICUBIC,
-                              "lanczos": PIL.Image.LANCZOS,
+        self.interpolation = {"linear": transforms.InterpolationMode.NEAREST,
+                              "bilinear": transforms.InterpolationMode.BILINEAR,
+                              "bicubic": transforms.InterpolationMode.BICUBIC,
+                              "lanczos": transforms.InterpolationMode.LANCZOS,
                               }[interpolation]
         self.image_transforms = transforms.Compose(
             [
-                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.Resize(size, interpolation=self.interpolation),
                 transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
                 transforms.RandomHorizontalFlip(),
                 transforms.ToTensor(),