import sys
import torch
from torch import nn

####################################################################################################

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

        self.layer_1 = nn.Sequential()
        self.layer_1.add_module("BN", nn.BatchNorm1d(num_features=100, eps=1e-5))
        self.layer_1.add_module("FC", nn.Linear(in_features=100, out_features=128))
        self.layer_1.add_module("AN", nn.ReLU())

        self.layer_2 = nn.Sequential()
        self.layer_2.add_module("BN", nn.BatchNorm1d(num_features=128, eps=1e-5))
        self.layer_2.add_module("FC", nn.Linear(in_features=128, out_features=1))
    
    def forward(self, x1, x2=None):
        if x2 is not None:
            x1 = self.layer_1(x1)
            x2 = self.layer_2(x2)
            x = torch.abs(x1 - x2)
            return self.layer_2(x)
        else:
            return self.layer_2(self.layer_1(x1))

####################################################################################################

def main():
    torch.manual_seed(0)
    model = Model()
    dummy = torch.randn(4, 100)

    output = model(dummy)

    print("\n" + 50*"*" + "\n")
    print("Output Shape : {}".format(output.shape))
    print("\n" + 50*"*" + "\n")

####################################################################################################

if __name__ == "__main__":
    sys.exit(main() or 0)

####################################################################################################

import sys
import torch
from torch import nn

####################################################################################################

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

        self.layer_1 = nn.Sequential()
        self.layer_1.add_module("BN", nn.BatchNorm1d(num_features=100, eps=1e-5))
        self.layer_1.add_module("FC", nn.Linear(in_features=100, out_features=128))
        self.layer_1.add_module("AN", nn.ReLU())

        self.layer_2 = nn.Sequential()
        self.layer_2.add_module("BN", nn.BatchNorm1d(num_features=128, eps=1e-5))
        self.layer_2.add_module("FC", nn.Linear(in_features=128, out_features=1))
    
    def forward(self, x1, x2=None):
        if x2 is not None:
            x1 = self.layer_1(x1)
            x2 = self.layer_2(x2)
            x = torch.abs(x1 - x2)
            return self.layer_2(x)
        else:
            return self.layer_2(self.layer_1(x1))

####################################################################################################

def main():
    torch.manual_seed(0)
    model = Model()
    dummy = torch.randn(4, 100)

    output = model(dummy)

    print("\n" + 50*"*" + "\n")
    print("Output Shape : {}".format(output.shape))
    print("\n" + 50*"*" + "\n")

####################################################################################################

if __name__ == "__main__":
    sys.exit(main() or 0)

####################################################################################################

#################################################################################################### #################################################################################################### #################################################################################################### ####################################################################################################