Commit aaf591bc authored by Chris Jewell's avatar Chris Jewell
Browse files

Moved output object from GPU onto host.

parent 996a4afc
......@@ -38,5 +38,6 @@ def chain_binomial_simulate(hazard_fn, state, start, end, time_step, stoichiomet
state = propagate(state)
output = output.write(i, state)
sim = output.gather(tf.range(times.shape[0]))
with tf.device("/CPU:0"):
sim = output.gather(tf.range(times.shape[0]))
return times, sim
......@@ -52,7 +52,7 @@ class HHstochastic(CovidUK):
if __name__=='__main__':
hh = draw_households(100, hh_size_distr)
K = (hh[:, None] == hh[None, :])
nsim = 120
nsim = 1000
model = HHstochastic(hh.astype(np.float64))
init_state = np.stack([np.broadcast_to(hh, [nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]]),
......
Supports Markdown
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