network.py

import torch.nn as nn

class Net(nn.Module):
  def __init__(self,input_size,hidden_size1,hidden_size2,output_size):
    super(Net, self).__init__()
    self.layer1=nn.Linear(input_size,hidden_size1)
    self.layer2=nn.Linear(hidden_size1,hidden_size2)
    self.layer3=nn.Linear(hidden_size2,output_size)
    self.relu=nn.ReLU()
  def forward(self,x):
    x=x.flatten(start_dim=1)
    x=self.layer1(x)
    x=self.relu(x)
    x=self.layer2(x)
    x=self.relu(x)
    x=self.layer3(x)

    return x