Tensor must be from the same graph

Hi,
I am having problems sampling use the most simplest model:
x <- normal(0,1)
m <- model(x)
mcmc(m)

My environment/installation:
jupyther data science notebook running in docker
installed greta according to the getting started tutorial.

Getting this error message:

Error in py_call_impl(callable, dots$args, dots$keywords): RuntimeError: Evaluation error: ValueError: Tensor(“leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/strided_slice:0”, shape=(), dtype=int32) must be from the same graph as Tensor(“mcmc_sample_chain/trace_scan/while/Const:0”, shape=(1,), dtype=float64).

Detailed traceback:
File “/home/jovyan/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py”, line 11055, in tile
“Tile”, input=input, multiples=multiples, name=name)
File “/home/jovyan/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py”, line 309, in _apply_op_helper
g = ops._get_graph_from_inputs(_Flatten(keywords.values()))
File “/home/jovyan/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py”, line 5921, in _get_graph_from_inputs
_assert_same_graph(original_graph_element, graph_element)
File “/home/jovyan/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py”, line 5856, in _assert_same_graph
(item, original_item))

Traceback:

  1. mcmc(m)
  2. lapply(initial_values_split, build_sampler, sampler, model)
  3. FUN(X[[i]], …)
  4. sampler$class$new(initial_values, model, sampler$parameters,
    . seed = seed)
  5. .subset2(public_bind_env, “initialize”)(…)
  6. self$define_tf_draws()
  7. dag$tf_run(sampler_batch <- tfp$mcmc$sample_chain(num_results = tf$math$floordiv(sampler_burst_length,
    . sampler_thin), current_state = free_state, kernel = sampler_kernel,
    . trace_fn = function(current_state, kernel_results) {
    . kernel_results
    . }, num_burnin_steps = tf$constant(0L, dtype = tf$int32),
    . num_steps_between_results = sampler_thin, parallel_iterations = 1L))
  8. self$on_graph(with(tfe, eval(expr)))
  9. with(self$tf_graph$as_default(), expr)
  10. with.python.builtin.object(self$tf_graph$as_default(), expr)
  11. tryCatch(force(expr), finally = {
    . data$__exit__(NULL, NULL, NULL)
    . })
  12. tryCatchList(expr, classes, parentenv, handlers)
  13. force(expr)
  14. with(tfe, eval(expr))
  15. with.default(tfe, eval(expr))
  16. eval(substitute(expr), data, enclos = parent.frame())
  17. eval(substitute(expr), data, enclos = parent.frame())
  18. eval(expr)
  19. eval(expr)
  20. tfp$mcmc$sample_chain(num_results = tf$math$floordiv(sampler_burst_length,
    . sampler_thin), current_state = free_state, kernel = sampler_kernel,
    . trace_fn = function(current_state, kernel_results) {
    . kernel_results
    . }, num_burnin_steps = tf$constant(0L, dtype = tf$int32),
    . num_steps_between_results = sampler_thin, parallel_iterations = 1L)
  21. py_call_impl(callable, dots$args, dots$keywords)

According to the getting started tutorial TF 1.14 or above should be installed, I interpreted this to also mean 2.0 but now I re-initialized everything with exactly TF 1.14.0 and it works.