diff --git a/docs/examples/2d/02-train.py b/docs/examples/2d/02-train.py
index 66170ea..619ca0d 100644
--- a/docs/examples/2d/02-train.py
+++ b/docs/examples/2d/02-train.py
@@ -41,7 +41,7 @@
 # We set the `max_iterations` equal to 5000 for demonstration purposes.
 # <br>(This takes around 20 minutes on a Mac Book Pro with an Apple M2 Max chip).
 
-device = "mps"  # 'mps', 'cpu', 'cuda:0'
+device = "cuda:0"  # 'mps', 'cpu', 'cuda:0'
 max_iterations = 5000
 
 train_config = TrainConfig(
@@ -62,5 +62,6 @@
 # Now we can begin the training! <br>
 # Uncomment the next two lines to train the model.
 
+# +
 # from cellulus.train import train
 # train(experiment_config)
diff --git a/docs/examples/2d/03-infer.py b/docs/examples/2d/03-infer.py
index ca2ad3d..1e9d20e 100644
--- a/docs/examples/2d/03-infer.py
+++ b/docs/examples/2d/03-infer.py
@@ -82,7 +82,7 @@
 # <br> The device could be set equal to `cuda:n` (where `n` is the index of
 # the GPU, for e.g. `cuda:0`), `cpu` or `mps`.
 
-device = "mps"  # "cuda:0", 'mps', 'cpu'
+device = "cuda:0"  # "cuda:0", 'mps', 'cpu'
 
 # We initialize the `inference_config` which contains our
 # `embeddings_dataset_config`, `segmentation_dataset_config` and
@@ -91,7 +91,6 @@
 # would like the cell membrane to be segmented or the nucleus.
 
 post_processing = "nucleus"
-bandwidth = 15.0
 
 inference_config = InferenceConfig(
     dataset_config=asdict(dataset_config),
@@ -100,7 +99,6 @@
     post_processed_dataset_config=asdict(post_processed_dataset_config),
     post_processing=post_processing,
     device=device,
-    bandwidth=bandwidth,
 )
 
 # ## Initialize `experiment_config`
@@ -180,4 +178,3 @@
     bottom_left_cmap=new_cmp,
     bottom_right_cmap=new_cmp,
 )
-# -