fix wrap on wm

Author: Vlad-Kor

model/distributions/torus/wrapped_normal/frolov.py

@@ -39,4 +39,5 @@ def sample(self, sample_options, distribution_options):
 
 		# wrapp
 		gaus_grid[:,0] = gaus_grid[:,0] % (2 * np.pi)
+		gaus_grid[:,1] = gaus_grid[:,1] % (2 * np.pi)
 		return gaus_grid

model/distributions/torus/wrapped_normal/improved_frolov.py

@@ -39,4 +39,5 @@ def sample(self, sample_options, distribution_options):
 
 		# wrapp
 		gaus_grid[:,0] = gaus_grid[:,0] % (2 * np.pi)
+		gaus_grid[:,1] = gaus_grid[:,1] % (2 * np.pi)
 		return gaus_grid