File size: 410 Bytes
70a6fd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class FeedForward(nn.Module):
def __init__(self, d_model, ff_dim=2048):
super(FeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, ff_dim)
self.linear2 = nn.Linear(ff_dim, d_model)
def forward(self, x):
x = F.relu(self.linear1(x))
x = self.linear2(x)
return x
|