Hello! I am having some issues on using batch norm. I am in the beginning of building my NN so for now I am using 100 samples for training and I want to overfit to it, just to make sure that the network can learn. The input and output are bout 1500 each. Here is my network:

```
class GW_NN(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(n_inp, 2000)
self.linear2 = nn.Linear(2000, 2000)
self.linear3 = nn.Linear(2000, 2000)
self.linear4 = nn.Linear(2000, 2000)
self.linear5 = nn.Linear(2000,n_out)
self.bn = nn.BatchNorm1d(2000)
def forward(self, x):
x = F.softplus(self.bn(self.linear1(x)))
x = F.softplus(self.bn(self.linear2(x)))
x = F.softplus(self.bn(self.linear3(x)))
x = F.softplus(self.bn(self.linear4(x)))
x = self.linear5(x)
return x
model_gw = GW_NN().cuda()
lrs = 1e-2
optimizer_gw = optim.Adam(model_gw.parameters(), lr = lrs)
for epoch in range(10001):
model_gw.train()
for i, dtt in enumerate(my_dataloader):
optimizer_gw.zero_grad()
inp = dtt[0].float().cuda()
output = dtt[1].float().cuda()
loss = F.mse_loss(model_gw(inp),output)
loss.backward()
optimizer_gw.step()
if epoch%100==0:
print(loss.data.cpu().numpy())
```

The loss goes down okish:

```
29418.57
20.279129
11.549426
8.563468
8.235117
8.161551
9.561671
7.5749683
7.60303
7.609553
7.265949
7.824227
10.810941
7.803124
7.6215243
7.977992
7.9355087
7.574047
7.326716
```

But when I want to try the trained NN (on the same data used for training) it fails:

```
idx = 10
y_real = output_data_phi[idx].data.cpu().numpy()
model_gw.eval()
y_pred = model_gw(input_data)[idx].data.cpu().numpy()
print(((y_pred-y_real)**2).mean())
```

I am getting `169114.33`

. I assume that the problem is with using batch norm in eval mode, but I am not sure what to do. Can someone help me? Thank you!