class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout = 0.3, max_seq_len = 200, device = 'cpu'):
super().__init__()
self.d_model = d_model
self.dropout = nn.Dropout(dropout)
pe = torch.zeros(max_seq_len, d_model).to(device)
pos = torch.arange(0, max_seq_len).unsqueeze(1).float()
two_i = torch.arange(0, d_model, step=2).float()
denominator = torch.pow(10000, two_i)
pe[:, 0::2] = torch.sin(pos/ denominator)
pe[:, 1::2] = torch.cos(pos/ denominator)
pe = pe.unsqueeze(0)
# assigns the first argument to a class variable
# i.e. self.pe
self.register_buffer("pe", pe)