fixes slurm weights saving

tests/trainer/test_trainer_steps.py

@@ -23,12 +23,11 @@ def test_trainingstep_dict(tmpdir):
         break
 
     out = trainer.run_training_batch(batch, batch_idx)
-    signal, grad_norm_dic, all_log_metrics, training_step_output_for_epoch_end = out
-    assert signal == 0
-    assert all_log_metrics['log_acc1'] == 12.0
-    assert all_log_metrics['log_acc2'] == 7.0
+    assert out.signal == 0
+    assert out.all_log_metrics['log_acc1'] == 12.0
+    assert out.all_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = training_step_output_for_epoch_end['pbar_on_batch_end']
+    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0
 
@@ -55,12 +54,11 @@ def training_step_with_step_end(tmpdir):
         break
 
     out = trainer.run_training_batch(batch, batch_idx)
-    signal, grad_norm_dic, all_log_metrics, training_step_output_for_epoch_end = out
-    assert signal == 0
-    assert all_log_metrics['log_acc1'] == 12.0
-    assert all_log_metrics['log_acc2'] == 7.0
+    assert out.signal == 0
+    assert out.all_log_metrics['log_acc1'] == 12.0
+    assert out.all_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = training_step_output_for_epoch_end['pbar_on_batch_end']
+    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0
 
@@ -92,12 +90,11 @@ def test_full_training_loop_dict(tmpdir):
         break
 
     out = trainer.run_training_batch(batch, batch_idx)
-    signal, grad_norm_dic, all_log_metrics, training_step_output_for_epoch_end = out
-    assert signal == 0
-    assert all_log_metrics['log_acc1'] == 12.0
-    assert all_log_metrics['log_acc2'] == 7.0
+    assert out.signal == 0
+    assert out.all_log_metrics['log_acc1'] == 12.0
+    assert out.all_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = training_step_output_for_epoch_end['pbar_on_batch_end']
+    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0
 
@@ -129,11 +126,10 @@ def test_train_step_epoch_end(tmpdir):
         break
 
     out = trainer.run_training_batch(batch, batch_idx)
-    signal, grad_norm_dic, all_log_metrics, training_step_output_for_epoch_end = out
-    assert signal == 0
-    assert all_log_metrics['log_acc1'] == 12.0
-    assert all_log_metrics['log_acc2'] == 7.0
+    assert out.signal == 0
+    assert out.all_log_metrics['log_acc1'] == 12.0
+    assert out.all_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = training_step_output_for_epoch_end['pbar_on_batch_end']
+    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0