title: tensorflow 2.0: An op outside of the function building code is being passed
body: <p>I'm getting an error: </p>
  
  <pre class="lang-py prettyprint-override"><code>TypeError: An op outside of the function building code is being passed
  a "Graph" tensor. It is possible to have Graph tensors
  leak out of the function building context by including a
  tf.init_scope in your function building code.
  For example, the following function will fail:
    @tf.function
    def has_init_scope():
      my_constant = tf.constant(1.)
      with tf.init_scope():
        added = my_constant * 2
  </code></pre>
  
  <p>Using an NVP layer like follows:</p>
  
  <pre class="lang-py prettyprint-override"><code>import tensorflow_probability as tfp
  tfb = tfp.bijectors
  tfd = tfp.distributions
  class NVPLayer(tf.keras.models.Model):
  
      def __init__(self, *, output_dim, num_masked, **kwargs):
          super().__init__(**kwargs)
          self.output_dim = output_dim
          self.num_masked = num_masked
          self.shift_and_log_scale_fn = tfb.real_nvp_default_template(
              hidden_layers=[2], # HERE HERE ADJUST THIS
              activation=None, # linear
              )
          self.loss = None
  
      def get_nvp(self):
          nvp = tfd.TransformedDistribution(
              distribution=tfd.MultivariateNormalDiag(loc=[0.] * self.output_dim),
              bijector=tfb.RealNVP(
                  num_masked=self.num_masked,
                  shift_and_log_scale_fn=self.shift_and_log_scale_fn)
              )
          return nvp
  
      def call(self, *inputs):
          nvp = self.get_nvp()
          self.loss = tf.reduce_mean(nvp.log_prob(*inputs)) # how else to do this?
          # return nvp.bijector.forward(*inputs)
          return nvp.bijector.inverse(*inputs)
  </code></pre>
  
  <p>I'm not calling <code>tf.init_scope</code> anywhere. A simple version training a layer like seems to work. </p>
  
  <p>I will try to get more more granular trace but I suspect this is something to do with non-eager mode stuff. </p>
  
  <p><strong>UPDATE:</strong> so this is definitely coming from the <code>self.loss</code> inclusion in some gradient tape layer. What is the correct way of doing this?</p>