import random
import numpy as np
import torch as tr

seed = 4312130                      # <<<<<<<<<<<<<<<< Your UPM ID Goes Here
random.seed(seed)
np.random.seed(seed)
tr.manual_seed(seed)

class DummyNet(tr.nn.Module):
    
  def __init__(self):
    super().__init__() 
    self.linear1 = tr.nn.Linear(784, 128, bias=True)       # 784 -> 16 - Layer 1 -- Affine Transformation (Linear with Bias)
    self.linear2 = tr.nn.Linear(128, 64, bias=False)       # 16 -> 16  - Layer 2 -- Linear Transformation (no bias)
    self.linear3 = tr.nn.Linear(64, 32, bias=True)        # 16 -> 10  - Layer 3 -- Affine Transformation (Linear with Bias)
    self.linear4 = tr.nn.Linear(32, 10, bias=False)
    
    
    self.init_weights()

  def init_weights(self):
    tr.nn.init.normal_(self.linear1.weight, mean=0.0, std=0.01)
    tr.nn.init.normal_(self.linear2.weight, mean=0.0, std=0.01)
    tr.nn.init.normal_(self.linear3.weight, mean=0.0, std=0.01)
    tr.nn.init.normal_(self.linear4.weight, mean=0.0, std=0.01)

    tr.nn.init.zeros_(self.linear1.bias)
    tr.nn.init.zeros_(self.linear3.bias)

   

    
  
  def forward(self, x):
    x = self.linear1(x)
    x = self.linear2(x)
    x = self.linear3(x)
    x = self.linear4(x)
   
    return x
      
model = DummyNet()
optimizer = tr.optim.SGD(model.parameters(), lr=0.01, maximize=False)   # Optimizer


loss_fn = tr.nn.CrossEntropyLoss()                                      # Objective Function [DO NOT CHANGE !]