diff --git a/.github/workflows/build-site.yaml b/.github/workflows/build-site.yaml
index 711b3b2..59c7fd2 100644
--- a/.github/workflows/build-site.yaml
+++ b/.github/workflows/build-site.yaml
@@ -22,6 +22,9 @@ jobs:
         with:
           python-version: '3.8'
       ##############################################
+      - name: Update pip and setuptools.
+        run: pip install --upgrade pip setuptools==57.1.0 wheel
+
       - name: Install pyg special-like.
         run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
 
diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml
index 821a18d..8965179 100644
--- a/.github/workflows/run-tests.yaml
+++ b/.github/workflows/run-tests.yaml
@@ -32,6 +32,8 @@ jobs:
         with:
           python-version: '3.8'
       ##############################################
+      - name: Update pip and setuptools.
+        run: pip install --upgrade pip setuptools==57.1.0 wheel
 
       - name: Install pyg special-like.
         run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.13.0+cpu.html
diff --git a/flowbot3d/datasets/flow_dataset_pyg.py b/flowbot3d/datasets/flow_dataset_pyg.py
index 958240e..1958951 100644
--- a/flowbot3d/datasets/flow_dataset_pyg.py
+++ b/flowbot3d/datasets/flow_dataset_pyg.py
@@ -23,6 +23,7 @@ def __init__(
         randomize_joints: bool = True,
         randomize_camera: bool = True,
         n_points: Optional[int] = 1200,
+        seed: int = 42,
     ) -> None:
         super().__init__()
 
@@ -33,12 +34,13 @@ def __init__(
             randomize_camera,
             n_points,
         )
+        self.seed = seed
 
     def len(self) -> int:
         return len(self.dataset)
 
     def get(self, index) -> tgd.Data:
-        return self.get_data(self.dataset._dataset._ids[index], seed=None)
+        return self.get_data(self.dataset._dataset._ids[index])
 
     @staticmethod
     def get_processed_dir(randomize_joints, randomize_camera):
@@ -46,8 +48,8 @@ def get_processed_dir(randomize_joints, randomize_camera):
         camera_chunk = "rc" if randomize_camera else "sc"
         return f"processed_{joint_chunk}_{camera_chunk}"
 
-    def get_data(self, obj_id: str, seed=None) -> Flowbot3DTGData:
-        data_dict = self.dataset.get_data(obj_id, seed)
+    def get_data(self, obj_id: str) -> Flowbot3DTGData:
+        data_dict = self.dataset.get_data(obj_id, self.seed)
 
         data = tgd.Data(
             id=data_dict["id"],
diff --git a/flowbot3d/evaluate_flow_models.py b/flowbot3d/evaluate_flow_models.py
index cf93564..51f6f48 100644
--- a/flowbot3d/evaluate_flow_models.py
+++ b/flowbot3d/evaluate_flow_models.py
@@ -40,7 +40,7 @@ def make_eval_dataset(
             randomize_camera,
         ),
         n_repeat=1,
-        n_proc=n_proc,
+        n_proc_per_worker=n_proc,
         seed=42,
     )
 
diff --git a/flowbot3d/train.py b/flowbot3d/train.py
index 4c412f2..d42fdb4 100644
--- a/flowbot3d/train.py
+++ b/flowbot3d/train.py
@@ -42,6 +42,7 @@ def create_flowbot_datasets(
     dataset: str,
     n_proc=-1,
     randomize_camera: bool = True,
+    seed=42,
 ) -> Tuple[tgd.Dataset, tgd.Dataset, tgd.Dataset]:
     if dataset == "umpnet":
         train_dset = CachedByKeyDataset(
@@ -50,6 +51,7 @@ def create_flowbot_datasets(
                 root=root / "raw",
                 split="umpnet-train-train",
                 randomize_camera=randomize_camera,
+                seed=seed,
             ),
             data_keys=rpd.UMPNET_TRAIN_TRAIN_OBJ_IDS,
             root=root,
@@ -58,7 +60,8 @@ def create_flowbot_datasets(
                 randomize_camera,
             ),
             n_repeat=100,
-            n_proc=n_proc,
+            n_workers=32,
+            n_proc_per_worker=n_proc,
         )
 
         test_dset = CachedByKeyDataset(
@@ -67,6 +70,7 @@ def create_flowbot_datasets(
                 root=root / "raw",
                 split="umpnet-train-test",
                 randomize_camera=randomize_camera,
+                seed=seed,
             ),
             data_keys=rpd.UMPNET_TRAIN_TEST_OBJ_IDS,
             root=root,
@@ -75,7 +79,8 @@ def create_flowbot_datasets(
                 randomize_camera,
             ),
             n_repeat=1,
-            n_proc=n_proc,
+            n_workers=32,
+            n_proc_per_worker=n_proc,
         )
 
         unseen_dset = CachedByKeyDataset(
@@ -84,6 +89,7 @@ def create_flowbot_datasets(
                 root=root / "raw",
                 split="umpnet-test",
                 randomize_camera=randomize_camera,
+                seed=seed,
             ),
             data_keys=rpd.UMPNET_TEST_OBJ_IDS,
             root=root,
@@ -92,7 +98,8 @@ def create_flowbot_datasets(
                 randomize_camera,
             ),
             n_repeat=1,
-            n_proc=n_proc,
+            n_workers=32,
+            n_proc_per_worker=n_proc,
         )
     elif dataset == "single":
         dset = CachedByKeyDataset(
@@ -101,6 +108,7 @@ def create_flowbot_datasets(
                 root=root / "raw",
                 split=["7179"],
                 randomize_camera=randomize_camera,
+                seed=seed,
             ),
             data_keys=["7179"],
             root=root,
@@ -109,7 +117,8 @@ def create_flowbot_datasets(
                 randomize_camera,
             ),
             n_repeat=1,
-            n_proc=n_proc,
+            n_workers=32,
+            n_proc_per_worker=n_proc,
         )
         train_dset = dset
         test_dset = dset
@@ -257,6 +266,7 @@ def train(
             dataset,
             n_proc,
             randomize_camera=randomize_camera,
+            seed=seed,
         )
 
     train_loader = tgl.DataLoader(train_dset, batch_size, shuffle=True, num_workers=0)
@@ -276,6 +286,7 @@ def train(
         project = "flowbot3d"
         logger = plog.WandbLogger(
             project=project,
+            entity="r-pad",
             config={
                 "dataset": dataset,
                 "model": model_type,
@@ -306,7 +317,8 @@ def train(
 
     # Create the trainer, which we'll train on only 1 gpu.
     trainer = pl.Trainer(
-        gpus=1,
+        accelerator="gpu",
+        devices=[0],
         logger=logger,
         callbacks=cbs,
         log_every_n_steps=log_every_n_steps,
diff --git a/pyproject.toml b/pyproject.toml
index 60bba7f..8ce4fb0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,8 @@ authors = [
 ]
 dependencies = [
   "gif",
-  "gym <= 0.20.0",
+  # "gym <= 0.20.0",  # This specific version is probably desired, but the CI no longer seems to build :( will fix later.
+  "gym",
   "h5py",
   "numpy <= 1.23.1",
   "open3d",