Simplified memory bank for Emformer

egs/librispeech/ASR/conv_emformer_transducer_stateless/emformer.py

@@ -1252,6 +1252,10 @@ def __init__(
     ):
         super().__init__()
 
+        assert int(math.log(chunk_length, 2)) == math.log(
+            chunk_length, 2
+        ), "chunk_length should be a power of 2."
+
         self.use_memory = memory_size > 0
         self.init_memory_op = nn.AvgPool1d(
             kernel_size=chunk_length,
@@ -1580,10 +1584,8 @@ def infer(
         chunk_mask = make_pad_mask(output_lengths).to(x.device)
         memory_mask = (
             (
-                torch.div(
-                    num_processed_frames,
-                    self.chunk_length,
-                    rounding_mode="floor",
+                (
+                    num_processed_frames << int(math.log(self.chunk_length, 2))
                 ).view(x.size(1), 1)
                 <= torch.arange(self.memory_size, device=x.device).expand(
                     x.size(1), self.memory_size

egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer.py

@@ -1188,6 +1188,10 @@ def __init__(
     ):
         super().__init__()
 
+        assert int(math.log(chunk_length, 2)) == math.log(
+            chunk_length, 2
+        ), "chunk_length should be a power of 2."
+
         self.use_memory = memory_size > 0
 
         self.emformer_layers = nn.ModuleList(
@@ -1488,10 +1492,8 @@ def infer(
         chunk_mask = make_pad_mask(output_lengths).to(x.device)
         memory_mask = (
             (
-                torch.div(
-                    num_processed_frames,
-                    self.chunk_length,
-                    rounding_mode="floor",
+                (
+                    num_processed_frames << int(math.log(self.chunk_length, 2))
                 ).view(x.size(1), 1)
                 <= torch.arange(self.memory_size, device=x.device).expand(
                     x.size(1), self.memory_size