Update Util_funs.py

Util_funs.py

@@ -191,10 +191,10 @@ def meta_train(data, model, device, Info,
 
 
 
-def train_loop(data_train_loader, data_test_loader, model, device, epoch = 4, lr = 1, print_info = True, name = 'name'):
+def train_loop(data_train_loader, data_test_loader, model, device, epoch = 4, lr = 1, print_info = True, name = 'name', weight_decay = 1):
   # Start the model's parameters
   model_meta = deepcopy(model)
-  optimizer = Adam(model_meta.parameters(), lr=lr)
+  optimizer = Adam(model_meta.parameters(), lr=lr, weight_decay = weight_decay)
 
   model_meta.to(device)
   model_meta.train()