import torch.nn as nn

class DaNN(nn.Module):

def __init__(self, n_input=28 * 28, n_hidden=256, n_class=10):

super(DaNN, self).__init__()

self.layer_input = nn.Linear(n_input, n_hidden)

self.dropout = nn.Dropout(p=0.5)

self.relu = nn.ReLU()

self.layer_hidden = nn.Linear(n_hidden, n_class)

def forward(self, src, tar):

x_src = self.layer_input(src)

x_tar = self.layer_input(tar)

x_src = self.dropout(x_src)

x_tar = self.dropout(x_tar)

x_src_mmd = self.relu(x_src)

x_tar_mmd = self.relu(x_tar)

y_src = self.layer_hidden(x_src_mmd)

return y_src, x_src_mmd, x_tar_mmd