Commit b43228de authored by Reimar H Leike's avatar Reimar H Leike
Browse files

fixed randomom state desync when using mpi kl

parent 116a16c6
Pipeline #63007 passed with stages
in 7 minutes and 10 seconds
...@@ -161,6 +161,7 @@ class MetricGaussianKL_MPI(Energy): ...@@ -161,6 +161,7 @@ class MetricGaussianKL_MPI(Energy):
if napprox > 1: if napprox > 1:
met._approximation = makeOp(approximation2endo(met, napprox)) met._approximation = makeOp(approximation2endo(met, napprox))
_samples = [] _samples = []
rand_state = np.random.get_state()
for i in range(lo, hi): for i in range(lo, hi):
if mirror_samples: if mirror_samples:
np.random.seed(i//2+seed_offset) np.random.seed(i//2+seed_offset)
...@@ -173,6 +174,7 @@ class MetricGaussianKL_MPI(Energy): ...@@ -173,6 +174,7 @@ class MetricGaussianKL_MPI(Energy):
else: else:
np.random.seed(i) np.random.seed(i)
_samples.append(met.draw_sample(from_inverse=True)) _samples.append(met.draw_sample(from_inverse=True))
np.random.set_state(rand_state)
_samples = tuple(_samples) _samples = tuple(_samples)
if mirror_samples: if mirror_samples:
n_samples *= 2 n_samples *= 2
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment