ml_pipeline/model/linear.py

11 lines
228 B
Python
Raw Normal View History

2022-10-29 17:19:59 -07:00
from torch import nn
class DNN(nn.Module):
def __init__(self, in_dim, out_dim):
super(DNN, self).__init__()
self.layer1 = nn.Linear(in_dim, out_dim)
def forward(self, x):
return self.layer1(x)