This is my model definition. You can choose between the dropout implementation. Thank you so much for your help

```
def __init__(self):
super(Network, self).__init__()
self.conv1 = nn.Conv2d(1,8,kernel_size = (5,5))
self.conv2 = nn.Conv2d(8,16, kernel_size = (3,3))
self.conv3 = nn.Conv2d(16, 32, kernel_size = (3,3))
self.conv4 = nn.Conv2d(32,64, kernel_size= (5,5))
#Maxpool Layers
self.maxp1 = nn.MaxPool2d((2,2))
self.maxp2 = nn.MaxPool2d((2,2))
self.maxp3 = nn.MaxPool2d((2,2))
self.maxp4 = nn.MaxPool2d((2,2))
#Dropout Layer
self.dropout = nn.Dropout(0.25)
#Linear Layers
self.linear1 = nn.Linear(64*5*5, 600)
self.linear2 = nn.Linear(600,100)
self.linear3 = nn.Linear(100,2)
@ torch.jit.script_method
def forward(self,x):
x = F.relu(self.conv1(x))# In: 1*128*128, Out: 8x124x124
x = self.maxp1(x) # In: 8x124x124, Out: 8*62*62
#x = self.dropout(x)
#x = F.dropout(x)
#print(x.size())
x = F.relu(self.conv2(x))# In: 8*62*62, Out: 16*60*60
print(x.size())
x = self.maxp2(x) # In: 16x60x60, Out: 16*30*30
#x = self.dropout(x)
#x = F.dropout(x)
x = F.relu(self.conv3(x))# In: 16*30*30, Out: 32*28*28
#print(x.size())
x = self.maxp3(x) # In: 32*28*28, Out: 32*14*14
#x = self.dropout(x)
#x = F.dropout(x)
x = F.relu(self.conv4(x))# In: 32*14*14, Out:64*10*10
#print(x.size())
x = self.maxp4(x)
# In: 64*10*10, Out:64*5*5
#x = F.dropout(x)
#x = self.dropout(x)
x = x.view(-1,64*5*5)
x = F.relu(self.linear1(x))
#x = F.dropout(x)
#x = self.dropout(x)
x = F.relu(self.linear2(x))
#x = F.dropout(x)
#x = self.dropout(x)
x =self.linear3(x)
x = F.softmax(x)
return x
```

Here ist the part, where most of the magic happens. I reduced it to the relevant code.

```
def train(network = None, epochs = 10, train_loader = None, learning_rate = 0.001):
network = network.train()
optimizer = opt.Adam(network.parameters(), lr=learning_rate)
loss_fn = torch.nn.CrossEntropyLoss()
current_acc = 0.0
for epoch in range(epochs):
for inputs, targets in train_loader:
optimizer.zero_grad()
outputs = network(inputs)
loss = loss_fn(outputs, targets)
loss.backward()
optimizer.step()
print("Epoch: ", epoch)
accuracy = evaluate(trained_network = network, eval_loader=eval_loader)
arrEpochs.append(epoch)
arrAccs.append(accuracy)
if accuracy > current_acc:
print("saving Model...")
current_acc = accuracy
network = network.eval()
torch.save(network.state_dict(), 'trained_model.pt')
network = network.train()
return network
```

When I push data manually through this ann I always set the network to eval()