Update modeling_conformer.py

modeling_conformer.py

@@ -41,7 +41,7 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
             else:
                 l.conv = nn.Sequential(nn.Conv2d(l.conv.in_channels,l.conv.out_channels,l.conv.kernel_size[0],l.conv.stride,1,groups=l.conv.out_channels),nn.Conv2d(l.conv.in_channels,l.conv.out_channels, 1))
         self.feature_extractor.conv_layers.append(Op(lambda x : x.transpose(1, 2)))
-        self.feature_projection.projection = nn.Linear(config.conv_dim[-1] * int(self.calc_length(torch.tensor(80.),repeat_num=config.num_feat_extract_layers)),config.hidden_size)
+        self.feature_projection.projection = nn.Linear(config.conv_dim[-1] * self.calc_length(80,repeat_num=config.num_feat_extract_layers),config.hidden_size)
         self.feature_projection.layer_norm = Op(lambda x:x.permute(0, 2, 1, 3).flatten(2))
         for l in self.encoder.layers:
             l.conv_module.glu = nn.Sequential(l.conv_module.glu,self.mask_layer)
@@ -61,10 +61,10 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         self.mask_layer.cache_pad_mask = (torch.arange(hidden_states.size(1), device=hidden_states.device).unsqueeze(0) >= self.cache_length.unsqueeze(1))
         return super()._mask_hidden_states(hidden_states, mask_time_indices, attention_mask)
 
-    def calc_length(self,lengths, all_paddings=2, kernel_size=3, stride=2, repeat_num=1):
+    def calc_length(self, lengths, all_paddings=2, kernel_size=3, stride=2, repeat_num=1):
         add_pad = all_paddings - kernel_size
         for _ in range(repeat_num):
-            lengths = torch.floor((lengths.float() + add_pad) / stride + 1)
+            lengths = (lengths + add_pad) // stride + 1
         return lengths
 
     def preprocessing(self, x):