Clean up save_1_or_2_dim_dualnet_visualizations interface

lnets/tasks/dualnets/mains/train_dual.py

@@ -106,7 +106,7 @@ def train_dualnet(model, loaders, config):
 
             # Visualize the learned critic landscape.
             if config.visualize:
-                save_1_or_2_dim_dualnet_visualizations(model, config.distrib1.dim, dirs.figures_dir, config,
+                save_1_or_2_dim_dualnet_visualizations(model, dirs.figures_dir, config,
                                                        state['epoch'], state['loss'])
 
         # Check if this is the best model.
@@ -158,7 +158,7 @@ def train_dualnet(model, loaders, config):
 
     # Visualize the learned critic landscape.
     if config.visualize:
-        save_1_or_2_dim_dualnet_visualizations(model, config.distrib1.dim, dirs.figures_dir, config,
+        save_1_or_2_dim_dualnet_visualizations(model, dirs.figures_dir, config,
                                                after_training=False)
 
     return test_state

lnets/utils/saving_and_loading.py

@@ -94,17 +94,19 @@ def load_best_model_and_optimizer(model, optimizer, best_path):
     load_optimizer(optimizer, best_optimizer_path)
 
 
-def save_1_or_2_dim_dualnet_visualizations(model, dim, figures_dir, config, epoch=None, loss=None,
+def save_1_or_2_dim_dualnet_visualizations(model, figures_dir, config, epoch=None, loss=None,
                                            after_training=False):
+
+    dim = config.distrib1.dim
     if not after_training:
         if dim == 2:
             save_2d_dualnet_visualizations(model, figures_dir, config, epoch, loss)
         if dim == 1:
             save_1d_dualnet_visualizations(model, figures_dir, config, epoch, loss)
     else:
-        if config.distrib1.dim == 2:
+        if dim == 2:
             save_2d_dualnet_visualizations(model, figures_dir, config, after_training=True)
-        if config.distrib1.dim == 1:
+        if dim == 1:
             save_1d_dualnet_visualizations(model, figures_dir, config, after_training=True)