diff --git a/BLcourse2.3/02_two_dim.py b/BLcourse2.3/02_two_dim.py
index 348949586aba23911d52bbf6cbcb118608cca49c..a02f4c8f09f5c5b5638a7167b42f3131e96bf7c2 100644
--- a/BLcourse2.3/02_two_dim.py
+++ b/BLcourse2.3/02_two_dim.py
@@ -136,7 +136,9 @@ if use_noise:
     # noisy train data
     noise_std = 0.2
     noise_dist = torch.distributions.Normal(loc=0, scale=noise_std)
-    y_train = data_train.z + noise_dist.sample_n(len(data_train.z))
+    y_train = data_train.z + noise_dist.sample(
+        sample_shape=(len(data_train.z),)
+    )
 else:
     # noise-free train data
     noise_std = 0