diff --git a/generation/maisi/maisi_train_vae_tutorial.ipynb b/generation/maisi/maisi_train_vae_tutorial.ipynb
index 26343a6a9..780e215a5 100644
--- a/generation/maisi/maisi_train_vae_tutorial.ipynb
+++ b/generation/maisi/maisi_train_vae_tutorial.ipynb
@@ -326,7 +326,7 @@
     }
    ],
    "source": [
-    "config_file = \"./configs/config_maisi.json\"\n",
+    "config_file = \"./configs/config_maisi3d-rflow.json\"\n",
     "config_dict = json.load(open(config_file, \"r\"))\n",
     "for k, v in config_dict.items():\n",
     "    setattr(args, k, v)\n",
diff --git a/generation/maisi/scripts/augmentation.py b/generation/maisi/scripts/augmentation.py
index 64469403a..55f586594 100644
--- a/generation/maisi/scripts/augmentation.py
+++ b/generation/maisi/scripts/augmentation.py
@@ -60,7 +60,7 @@ def dilate3d(input_tensor, erosion=3):
     return output.squeeze(0).squeeze(0)
 
 
-def augmentation_tumor_bone(pt_nda, output_size, random_seed):
+def augmentation_tumor_bone(pt_nda, output_size, random_seed=None):
     volume = pt_nda.squeeze(0)
     real_l_volume_ = torch.zeros_like(volume)
     real_l_volume_[volume == 128] = 1
@@ -113,7 +113,7 @@ def augmentation_tumor_bone(pt_nda, output_size, random_seed):
     return pt_nda
 
 
-def augmentation_tumor_liver(pt_nda, output_size, random_seed):
+def augmentation_tumor_liver(pt_nda, output_size, random_seed=None):
     volume = pt_nda.squeeze(0)
     real_l_volume_ = torch.zeros_like(volume)
     real_l_volume_[volume == 1] = 1
@@ -163,7 +163,7 @@ def augmentation_tumor_liver(pt_nda, output_size, random_seed):
     return pt_nda
 
 
-def augmentation_tumor_lung(pt_nda, output_size, random_seed):
+def augmentation_tumor_lung(pt_nda, output_size, random_seed=None):
     volume = pt_nda.squeeze(0)
     real_l_volume_ = torch.zeros_like(volume)
     real_l_volume_[volume == 23] = 1
@@ -227,7 +227,7 @@ def augmentation_tumor_lung(pt_nda, output_size, random_seed):
     return pt_nda
 
 
-def augmentation_tumor_pancreas(pt_nda, output_size, random_seed):
+def augmentation_tumor_pancreas(pt_nda, output_size, random_seed=None):
     volume = pt_nda.squeeze(0)
     real_l_volume_ = torch.zeros_like(volume)
     real_l_volume_[volume == 4] = 1
@@ -277,7 +277,7 @@ def augmentation_tumor_pancreas(pt_nda, output_size, random_seed):
     return pt_nda
 
 
-def augmentation_tumor_colon(pt_nda, output_size, random_seed):
+def augmentation_tumor_colon(pt_nda, output_size, random_seed=None):
     volume = pt_nda.squeeze(0)
     real_l_volume_ = torch.zeros_like(volume)
     real_l_volume_[volume == 27] = 1
@@ -335,7 +335,7 @@ def augmentation_tumor_colon(pt_nda, output_size, random_seed):
     return pt_nda
 
 
-def augmentation_body(pt_nda, random_seed):
+def augmentation_body(pt_nda, random_seed=None):
     volume = pt_nda.squeeze(0)
 
     zoom = RandZoom(min_zoom=0.99, max_zoom=1.01, mode="nearest", align_corners=None, prob=1.0)
@@ -347,7 +347,7 @@ def augmentation_body(pt_nda, random_seed):
     return pt_nda
 
 
-def augmentation(pt_nda, output_size, random_seed):
+def augmentation(pt_nda, output_size, random_seed=None):
     label_list = torch.unique(pt_nda)
     label_list = list(label_list.cpu().numpy())