diff --git a/src/models.py b/src/models.py
index 8e95fc5ec9f8a18627fc1703e591abfd605f16d0..cd01cbce2240fbb05229cfe25e862c9f0b4a313f 100644
--- a/src/models.py
+++ b/src/models.py
@@ -91,50 +91,4 @@ class GenericConvModel(ImageClassificationBase):
         )
         
     def forward(self, xb):
-        return self.network(xb)    
-        
-
-class CovNet(ImageClassificationBase):
-    def __init__(self, in_size, n_channels):
-        super().__init__()
-        self.img_final_size = int(in_size / (2**3))
-        self.network = nn.Sequential(
-            nn.Conv2d(n_channels, 32, kernel_size=3, padding=1),
-            nn.BatchNorm2d(32),
-            nn.ReLU(),
-            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(32),
-            nn.ReLU(),
-            nn.MaxPool2d(2, 2),  # output: 32 x 16 x 16
-            nn.Dropout(0.25),
-            
-            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(64),
-            nn.ReLU(),
-            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(64),
-            nn.ReLU(),
-            nn.MaxPool2d(2, 2),  # output: 64 x 8 x 8
-            nn.Dropout(0.25),
-            
-            
-            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(128),
-            nn.ReLU(),
-            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(128),
-            nn.ReLU(),
-            nn.MaxPool2d(2, 2),  # output: 128 x 4 x 4
-            nn.Dropout(0.25),
-            
-            nn.Flatten(),
-            nn.Linear(128 * self.img_final_size * self.img_final_size, 512),
-            nn.ReLU(),
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.Linear(256, 10)
-        )
-        
-    def forward(self, xb):
-        return self.network(xb)           
-        
+        return self.network(xb)   
\ No newline at end of file
diff --git a/src/utils_data.py b/src/utils_data.py
index d15f37bd86ba76aa26c9730079e2a32729f076d2..d28573784f369aa7cadac94503ef193f085f4076 100644
--- a/src/utils_data.py
+++ b/src/utils_data.py
@@ -53,7 +53,7 @@ def create_label_dict(dataset : str, nn_model : str) -> dict:
         fashion_mnist = torchvision.datasets.MNIST("datasets", download=True)
         (x_data, y_data) = fashion_mnist.data, fashion_mnist.targets
     
-        if nn_model in ["convolutional","CovNet"]:
+        if nn_model in ["convolutional"]:
             x_data = x_data.unsqueeze(1)
 
     elif dataset == 'mnist':
@@ -63,7 +63,7 @@ def create_label_dict(dataset : str, nn_model : str) -> dict:
     elif dataset == "cifar10":
         cifar10 = torchvision.datasets.CIFAR10("datasets", download=True)
         (x_data, y_data) = cifar10.data, cifar10.targets
-        #x_data = np.transpose(x_data, (0, 3, 1, 2))
+        
         
     elif dataset == 'kmnist':
         kmnist = torchvision.datasets.KMNIST("datasets", download=True)
@@ -256,9 +256,9 @@ def data_preparation(client: Client, row_exp: dict) -> None:
     test_dataset = CustomDataset(x_test, y_test, transform=test_val_transform)
 
     # Create DataLoaders
-    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-    validation_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)
-    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True)
+    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+    validation_loader = DataLoader(val_dataset, batch_size=128, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=True)
 
     # Store DataLoaders in the client object
     setattr(client, 'data_loader', {'train': train_loader, 'val': validation_loader, 'test': test_loader})
@@ -307,7 +307,7 @@ def setup_experiment(row_exp: dict) -> Tuple[Server, list]:
 
     """
 
-    from src.models import GenericConvModel, CovNet
+    from src.models import GenericConvModel
     from src.utils_fed import init_server_cluster
     import torch
     
@@ -326,8 +326,6 @@ def setup_experiment(row_exp: dict) -> Tuple[Server, list]:
     elif row_exp['nn_model'] == "convolutional": 
         
         model_server = Server(GenericConvModel(in_size=imgs_params[row_exp['dataset']][0], n_channels=imgs_params[row_exp['dataset']][1]))
-    elif row_exp['nn_model'] == "CovNet":
-        model_server = Server(CovNet(in_size=imgs_params[row_exp['dataset']][0], n_channels=imgs_params[row_exp['dataset']][1]))
        
 
     model_server.model.to(device)
@@ -613,7 +611,7 @@ def swap_labels(labels : list, client : Client, heterogeneity_class : int) -> Cl
     return client
 
 
-
+'''
 def centralize_data(list_clients : list) -> Tuple[DataLoader, DataLoader]:
     """Centralize data of the federated learning setup for central model comparison
 
@@ -643,11 +641,12 @@ def centralize_data(list_clients : list) -> Tuple[DataLoader, DataLoader]:
     x_test_tensor = torch.tensor(x_test, dtype=torch.float32)
     y_test_tensor = torch.tensor(y_test, dtype=torch.long)
     
-    train_loader = DataLoader(TensorDataset(x_train_tensor, y_train_tensor), batch_size=64, shuffle=True)
-    val_loader = DataLoader(TensorDataset(x_val_tensor, y_val_tensor), batch_size=64, shuffle=True)
-    test_loader = DataLoader(TensorDataset(x_test_tensor, y_test_tensor), batch_size=64, shuffle=True)
+    train_loader = DataLoader(TensorDataset(x_train_tensor, y_train_tensor), batch_size=128, shuffle=True)
+    val_loader = DataLoader(TensorDataset(x_val_tensor, y_val_tensor), batch_size=128, shuffle=False)
+    test_loader = DataLoader(TensorDataset(x_test_tensor, y_test_tensor), batch_size=128, shuffle=False)
     
     return train_loader, val_loader, test_loader
+'''
 
 def centralize_data(list_clients: list) -> Tuple[DataLoader, DataLoader]:
     """Centralize data of the federated learning setup for central model comparison
@@ -696,9 +695,9 @@ def centralize_data(list_clients: list) -> Tuple[DataLoader, DataLoader]:
     test_dataset = CustomDataset(x_test, y_test, transform=test_val_transform)
 
     # Create DataLoaders
-    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)  # Validation typically not shuffled
-    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)  # Test data typically not shuffled
+    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)  # Validation typically not shuffled
+    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)  # Test data typically not shuffled
 
     return train_loader, val_loader, test_loader
 
diff --git a/src/utils_fed.py b/src/utils_fed.py
index 375b2af2ff69be880a9c3fc31879f3819a866718..5c08df9a4da540774d3162bcbb0b40b90f4a4521 100644
--- a/src/utils_fed.py
+++ b/src/utils_fed.py
@@ -199,7 +199,7 @@ def init_server_cluster(my_server : Server, list_clients : list, row_exp : dict,
         p_expert_opintion : Parameter to avoid completly random assignment if neeed (default to 0)
     """
     
-    from src.models import GenericLinearModel, GenericConvModel, CovNet
+    from src.models import GenericLinearModel, GenericConvModel
     import numpy as np
     import copy
 
diff --git a/src/utils_training.py b/src/utils_training.py
index d1cc6d03d6e94316f29a4538d7e2bbc6cfda779f..ca2e523a3553279c11a0f244402c83a2929e9dc8 100644
--- a/src/utils_training.py
+++ b/src/utils_training.py
@@ -212,7 +212,7 @@ def train_central(model: ImageClassificationBase, train_loader: DataLoader, val_
     # Move the model to the appropriate device
     model.to(device)
 
-    opt_func = torch.optim.Adam  # if row_exp['nn_model'] == "linear" else torch.optim.Adam
+    opt_func = torch.optim.SGD  # if row_exp['nn_model'] == "linear" else torch.optim.Adam
     lr = 0.001
     history = []
     optimizer = opt_func(model.parameters(), lr)