reduce tracing

diffusion_sampler.py

@@ -118,7 +118,7 @@ class DiffusionSampler(keras.Model):
         sqrt_one_minus_alpha_cum_prod = at_timestep(self.sqrt_one_minus_alphas_cum_prod)
         return sqrt_alpha_cum_prod * x_start + sqrt_one_minus_alpha_cum_prod * noise
 
-    @tf.function()
+    @tf.function(reduce_tracing=True)
     def generate_images(
         self,
         num_images: int,