Check size on transformer cache

browsermt · Jun 8, 2022 · ca0aeb2 · ca0aeb2
1 parent e88c1aa
commit ca0aeb2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/models/transformer.h b/src/models/transformer.h
@@ -281,7 +281,7 @@ class Transformer : public EncoderOrDecoderBase {
     // memoization propagation (short-term)
     if (cache                                                                          // if caching
         && cache_.count(prefix + "_keys") > 0                                          // and the keys expression has been seen
-        && cache_[prefix + "_keys"]->shape().elements() == keys->shape().elements()) { // and the underlying element size did not change
+        && cache_[prefix + "_keys"]->shape() == keys->shape()) {                       // and the underlying shape did not change
       kh = cache_[prefix + "_keys"];                                                   // then return cached tensor
     }
     else {
@@ -296,7 +296,7 @@ class Transformer : public EncoderOrDecoderBase {
     Expr vh;
     if (cache 
         && cache_.count(prefix + "_values") > 0 
-        && cache_[prefix + "_values"]->shape().elements() == values->shape().elements()) {
+        && cache_[prefix + "_values"]->shape() == values->shape()) {
       vh = cache_[prefix + "_values"];
     } else {
       auto Wv = graph_->param(prefix + "_Wv", {dimModel, dimModel}, inits::glorotUniform());