testing map location

tests/test_models.py

@@ -21,53 +21,6 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
-def test_early_stopping_cpu_model():
-    """
-    Test each of the trainer options
-    :return:
-    """
-
-    stopping = EarlyStopping()
-    trainer_options = dict(
-        early_stop_callback=stopping,
-        gradient_clip=1.0,
-        overfit_pct=0.20,
-        track_grad_norm=2,
-        print_nan_grads=True,
-        progress_bar=False,
-        experiment=get_exp(),
-        train_percent_check=0.1,
-        val_percent_check=0.1
-    )
-
-    model, hparams = get_model()
-    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
-    # test freeze on cpu
-    model.freeze()
-    model.unfreeze()
-
-
-def test_cpu_model_with_amp():
-    """
-    Make sure model trains on CPU
-    :return:
-    """
-
-    trainer_options = dict(
-        progress_bar=False,
-        experiment=get_exp(),
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.4,
-        use_amp=True
-    )
-
-    model, hparams = get_model()
-
-    with pytest.raises(MisconfigurationException):
-        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
 
 def test_amp_gpu_ddp_slurm_managed():
     """
@@ -123,7 +76,8 @@ def test_amp_gpu_ddp_slurm_managed():
     assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
 
     # test model loading
-    pretrained_model = load_model(exp, save_dir, True)
+    map_location = 'cuda:1'
+    pretrained_model = load_model(exp, save_dir, True, map_location)
 
     # test model preds
     run_prediction(model.test_dataloader, pretrained_model)
@@ -144,6 +98,54 @@ def test_amp_gpu_ddp_slurm_managed():
     clear_save_dir()
 
 
+def test_early_stopping_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    stopping = EarlyStopping()
+    trainer_options = dict(
+        early_stop_callback=stopping,
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        train_percent_check=0.1,
+        val_percent_check=0.1
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = get_model()
+
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
 def test_cpu_model():
     """
     Make sure model trains on CPU
@@ -433,7 +435,7 @@ def clear_save_dir():
         shutil.rmtree(save_dir)
 
 
-def load_model(exp, save_dir, on_gpu):
+def load_model(exp, save_dir, on_gpu, map_location=None):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -442,7 +444,10 @@ def load_model(exp, save_dir, on_gpu):
     checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
     weights_dir = os.path.join(save_dir, checkpoints[0])
 
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=on_gpu)
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+                                                             tags_csv=tags_path,
+                                                             on_gpu=on_gpu,
+                                                             map_location=map_location)
 
     assert trained_model is not None, 'loading model failed'