fixed return elements (batch_size was missing)

create_parameter_weights.py

@@ -83,7 +83,7 @@ def main():
     print("Computing mean and std.-dev. for parameters...")
     means = []
     squares = []
-    for init_batch, target_batch, forcing_batch in tqdm(loader):
+    for init_batch, target_batch, _, forcing_batch in tqdm(loader):
         batch = torch.cat(
             (init_batch, target_batch), dim=1
         )  # (N_batch, N_t, N_grid, d_features)
@@ -128,7 +128,7 @@ def main():
     diff_means = []
     diff_squares = []
     for batch_data in tqdm(loader_standard):
-        init_batch, target_batch, _ = batch_data
+        init_batch, target_batch, _ , _= batch_data
         batch = torch.cat(init_batch, target_batch, dim=1)
         # (N_batch, N_t', N_grid, d_features)
         batch_diffs = batch[:, 1:] - batch[:, :-1]