Improve performance (#67)

amansrivastava17 · Aug 24, 2021 · 8cd4087 · 8cd4087
1 parent fe4393e
commit 8cd4087
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/server/embedding_as_service/text/xlnet/models/train.py b/server/embedding_as_service/text/xlnet/models/train.py
@@ -204,10 +204,10 @@ def get_cache_fn(mem_len):
   def cache_fn(batch_size):
     mems = []
     if FLAGS.mem_len > 0:
-      for _ in range(FLAGS.n_layer):
-        zeros = tf.zeros(
+      zeros = tf.zeros(
             [mem_len, batch_size, FLAGS.d_model],
             dtype=tf_float)
+      for _ in range(FLAGS.n_layer):
         mems.append(zeros)
 
     return mems