mirror of https://github.com/hpcaitech/ColossalAI
18 lines
385 B
Python
18 lines
385 B
Python
|
import torch
|
||
|
|
import colossalai.nn as col_nn
|
||
|
|
|
||
|
|
|
||
|
|
class MLP(torch.nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, dim: int, layers: int):
|
||
|
|
super().__init__()
|
||
|
|
self.layers = torch.nn.ModuleList()
|
||
|
|
|
||
|
|
for _ in range(layers):
|
||
|
|
self.layers.append(col_nn.Linear(dim, dim))
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
for layer in self.layers:
|
||
|
|
x = layer(x)
|
||
|
|
return x
|