gpt2 / src /model /feed_forward.py

Upload folder using huggingface_hub

c21e887 verified 4 months ago

342 Bytes

import torch.nn as nn


class FeedForwardNetwork(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, 4 * embed_dim),
            nn.GELU(),
            nn.Linear(4 * embed_dim, embed_dim),
        )

    def forward(self, x):
        return self.net(x)

	import torch.nn as nn


	class FeedForwardNetwork(nn.Module):
	def __init__(self, embed_dim):
	super().__init__()
	self.net = nn.Sequential(
	nn.Linear(embed_dim, 4 * embed_dim),
	nn.GELU(),
	nn.Linear(4 * embed_dim, embed_dim),
	)

	def forward(self, x):
	return self.net(x)