Update modeling_conformer.py

modeling_conformer.py

@@ -239,15 +239,16 @@ encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv
         t , s = x
         start_token_segment = self.config.languages.index(self.language) * self.joint.out_features
         all_tokens, all_starts, all_ends = [], [], []
+        device = t[0].device
         for tokens, starts, (s, e) in zip(t,s, ts):
             tokens += start_token_segment
             starts += s
             all_tokens.append(tokens)
             all_starts.append(starts)
             all_ends.append(torch.cat([starts[1:], e[None]]))
-            all_tokens.append(torch.tensor([-1]))
-            all_starts.append(torch.tensor([e]))
-            all_ends.append(torch.tensor([e + 0.005]))
+            all_tokens.append(torch.tensor([-1],device=device))
+            all_starts.append(torch.tensor([e],device=device))
+            all_ends.append(torch.tensor([e + 0.005],device=device))
         return [srt.Subtitle(i,timedelta(seconds=float(st)),timedelta(seconds=float(en)),"<line>" if tok == -1 else self.config.vocab[int(tok)]) for i, (tok, st, en) in enumerate(zip(torch.cat(all_tokens), torch.cat(all_starts), torch.cat(all_ends)), 1)]