shuffle and drop last

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rectified-flow-pytorch"
-version = "0.0.11"
+version = "0.0.12"
 description = "Rectified Flow in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

rectified_flow_pytorch/rectified_flow.py

@@ -736,7 +736,7 @@ def __init__(
 
         self.model = rectified_flow
         self.optimizer = Adam(rectified_flow.parameters(), lr = learning_rate, **adam_kwargs)
-        self.dl = DataLoader(dataset, batch_size = batch_size)
+        self.dl = DataLoader(dataset, batch_size = batch_size, shuffle = True, drop_last = True)
 
         self.model, self.optimizer, self.dl = self.accelerator.prepare(self.model, self.optimizer, self.dl)