Spaces:
Runtime error
Runtime error
File size: 987 Bytes
7900c16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
import torch
import torch.nn as nn
class PosEmbedding(nn.Module):
"""
Learnable Position Embedding
"""
def __init__(self, args, _):
super(PosEmbedding, self).__init__()
if "speech" in args.embedding:
self.max_seq_length = max(args.max_seq_length, args.max_audio_frames)
else:
self.max_seq_length = args.max_seq_length
self.embedding = nn.Embedding(self.max_seq_length, args.emb_size)
def forward(self, _, seg):
"""
Args:
src: [batch_size x seq_length]
seg: [batch_size x seq_length]
Returns:
emb: [batch_size x seq_length x hidden_size]
"""
seq_length = seg.size(1)
batch_size = seg.size(0)
device = seg.device
pos_emb = self.embedding(
torch.arange(0, seq_length, device=device, dtype=torch.long)
.unsqueeze(0)
.repeat(batch_size, 1)
)
return pos_emb
|