fix test keep eigen vec

python/tffpy/experiments/exp_variance.py

@@ -65,7 +65,8 @@ class VarianceExperiment(SolveTffExperiment):
     def __init__(self, force_reset=False, suffix=''):
         SolveTffExperiment.__init__(self,
                                     force_reset=force_reset,
-                                    suffix='Variance' + suffix)
+                                    suffix='Variance' + suffix,
+                                    keep_eigenvectors=[])
 
     def display_results(self):
         res = self.load_results(array_type='xarray')

python/tffpy/experiments/tests/test_exp_solve_tff.py

@@ -71,7 +71,7 @@ class TestSolveTffExperiment(unittest.TestCase):
 
         light_exp = SolveTffExperiment.get_experiment(
             setting='light', force_reset=False)
-        for idt in range(light_exp.n_tasks):
+        for idt in light_exp.keep_eigenvectors:
             light_exp.plot_task(idt=idt, fontsize=16)
             plt.close('all')
         light_exp.plot_results()

python/tffpy/experiments/tests/test_exp_variance.py

@@ -71,7 +71,7 @@ class TestVarianceExperiment(unittest.TestCase):
 
         light_exp = VarianceExperiment.get_experiment(
             setting='light', force_reset=False)
-        for idt in range(light_exp.n_tasks):
+        for idt in light_exp.keep_eigenvectors:
             light_exp.plot_task(idt=idt, fontsize=16)
             plt.close('all')
         light_exp.plot_results()

python/tffpy/experiments/yafe_slurm.py

@@ -61,7 +61,7 @@ import sys
 
 def generate_slurm_script(script_file_path, xp_var_name, task_ids=None,
                           n_simultaneous_jobs=10, slurm_walltime='02:00:00',
-                          activate_env_command=None, use_gpu=False):
+                          activate_env_command=None, use_cpu_gpu='cpu'):
     """Generate a script to launch an experiment using Slurm.
 
     Tasks are divided into batches that are executed by oar jobs.
@@ -95,9 +95,8 @@ def generate_slurm_script(script_file_path, xp_var_name, task_ids=None,
         ``source activate some_conda_env`` when using conda.
         If ``activate_env_command`` is ``None``, no virtual environment is
         activated.
-    use_gpu : bool
-        Flag specifying if a gpu ressource is needed when running the
-        experiment. This has not been implemented yet.
+    use_cpu_gpu : {'all', 'cpu', 'gpu'}
+        Parameter to choose using CPU, GPU or both.
     """
     script_file_path = Path(script_file_path)
     script_dir = script_file_path.parent
@@ -127,10 +126,8 @@ def generate_slurm_script(script_file_path, xp_var_name, task_ids=None,
     script += '#SBATCH --output={}/stdout_%A_%a.slurm\n'.format(log_dir)
     script += '#SBATCH --error={}/stderr_%A_%a.slurm\n'.format(log_dir)
     script += '#SBATCH --time={}\n'.format(slurm_walltime)
-    # if use_gpu:
-    #     script += '#SBATCH -p gpu IS NOT NULL\n'
-    # else:
-    #     script += '#SBATCH -p gpu IS NULL\n'
+    if use_cpu_gpu in ('cpu', 'gpu'):
+        script += '#SBATCH --partition={}\n'.format(use_cpu_gpu)
 
     script += 'srun -N1 -n1 {}/run_$SLURM_ARRAY_TASK_ID.sh'.format(script_dir)